unique_toolkit 0.8.18__py3-none-any.whl → 0.8.19__py3-none-any.whl

unique_toolkit/history_manager/loop_token_reducer.py

@@ -1,18 +1,31 @@
-
 import json
 from logging import Logger
 from typing import Awaitable, Callable
 
-from pydantic import BaseModel
 import tiktoken
-from unique_toolkit._common.token.token_counting import num_token_for_language_model_messages
+from pydantic import BaseModel
+
+from unique_toolkit._common.token.token_counting import (
+    num_token_for_language_model_messages,
+)
 from unique_toolkit._common.validators import LMI
 from unique_toolkit.app.schemas import ChatEvent
 from unique_toolkit.chat.service import ChatService
 from unique_toolkit.content.schemas import ContentChunk
 from unique_toolkit.content.service import ContentService
-from unique_toolkit.history_manager.history_construction_with_contents import FileContentSerialization, get_full_history_with_contents
-from unique_toolkit.language_model.schemas import LanguageModelAssistantMessage, LanguageModelMessage, LanguageModelMessageRole, LanguageModelMessages, LanguageModelSystemMessage, LanguageModelToolMessage, LanguageModelUserMessage
+from unique_toolkit.history_manager.history_construction_with_contents import (
+    FileContentSerialization,
+    get_full_history_with_contents,
+)
+from unique_toolkit.language_model.schemas import (
+    LanguageModelAssistantMessage,
+    LanguageModelMessage,
+    LanguageModelMessageRole,
+    LanguageModelMessages,
+    LanguageModelSystemMessage,
+    LanguageModelToolMessage,
+    LanguageModelUserMessage,
+)
 from unique_toolkit.reference_manager.reference_manager import ReferenceManager
 
 
@@ -26,18 +39,16 @@ class SourceReductionResult(BaseModel):
         arbitrary_types_allowed = True
 
 
-class LoopTokenReducer():
-    
+class LoopTokenReducer:
     def __init__(
         self,
         logger: Logger,
         event: ChatEvent,
-        max_history_tokens:int,
+        max_history_tokens: int,
         has_uploaded_content_config: bool,
         reference_manager: ReferenceManager,
-        language_model: LMI
+        language_model: LMI,
     ):
-        
         self._max_history_tokens = max_history_tokens
         self._has_uploaded_content_config = has_uploaded_content_config
         self._logger = logger
@@ -48,29 +59,28 @@ class LoopTokenReducer():
         self._content_service = ContentService.from_event(event)
         self._user_message = event.payload.user_message
         self._chat_id = event.payload.chat_id
-    
 
     def _get_encoder(self, language_model: LMI) -> tiktoken.Encoding:
         name = language_model.encoder_name or "cl100k_base"
         return tiktoken.get_encoding(name)
 
-    async def get_history_for_model_call( self,
+    async def get_history_for_model_call(
+        self,
         original_user_message: str,
         rendered_user_message_string: str,
         rendered_system_message_string: str,
         loop_history: list[LanguageModelMessage],
-        remove_from_text: Callable[[str], Awaitable[str]]
-        ) -> LanguageModelMessages:
+        remove_from_text: Callable[[str], Awaitable[str]],
+    ) -> LanguageModelMessages:
         """Compose the system and user messages for the plan execution step, which is evaluating if any further tool calls are required."""
 
-
         messages = await self._construct_history(
             original_user_message,
             rendered_user_message_string,
             rendered_system_message_string,
             loop_history,
-            remove_from_text
-            )
+            remove_from_text,
+        )
 
         token_count = self._count_message_tokens(messages)
         self._log_token_usage(token_count)
@@ -79,11 +89,11 @@ class LoopTokenReducer():
             token_count_before_reduction = token_count
             loop_history = self._handle_token_limit_exceeded(loop_history)
             messages = await self._construct_history(
-              original_user_message,
-              rendered_user_message_string,
-              rendered_system_message_string,
-              loop_history,
-              remove_from_text
+                original_user_message,
+                rendered_user_message_string,
+                rendered_system_message_string,
+                loop_history,
+                remove_from_text,
             )
             token_count = self._count_message_tokens(messages)
             self._log_token_usage(token_count)
@@ -92,7 +102,7 @@ class LoopTokenReducer():
                 break
 
         return messages
-    
+
     def _exceeds_token_limit(self, token_count: int) -> bool:
         """Check if token count exceeds the maximum allowed limit and if at least one tool call has more than one source."""
         # At least one tool call should have more than one chunk as answer
@@ -105,13 +115,11 @@ class LoopTokenReducer():
         # include system_prompt and user question already
         # TODO: There is a problem if we exceed but only have one chunk per tool call
         exceeds_limit = (
-            token_count
-            > self._language_model.token_limits.token_limit_input
+            token_count > self._language_model.token_limits.token_limit_input
         )
 
         return has_multiple_chunks_for_a_tool_call and exceeds_limit
 
-
     def _count_message_tokens(self, messages: LanguageModelMessages) -> int:
         """Count tokens in messages using the configured encoding model."""
         return num_token_for_language_model_messages(
@@ -130,28 +138,34 @@ class LoopTokenReducer():
         rendered_user_message_string: str,
         rendered_system_message_string: str,
         loop_history: list[LanguageModelMessage],
-        remove_from_text: Callable[[str], Awaitable[str]]
+        remove_from_text: Callable[[str], Awaitable[str]],
     ) -> LanguageModelMessages:
         history_from_db = await self._get_history_from_db(remove_from_text)
-        history_from_db = self._replace_user_message(history_from_db, original_user_message, rendered_user_message_string)
-        system_message = LanguageModelSystemMessage(content=rendered_system_message_string)
+        history_from_db = self._replace_user_message(
+            history_from_db, original_user_message, rendered_user_message_string
+        )
+        system_message = LanguageModelSystemMessage(
+            content=rendered_system_message_string
+        )
 
         constructed_history = LanguageModelMessages(
             [system_message] + history_from_db + loop_history,
         )
 
-
         return constructed_history
-    
 
-    def _handle_token_limit_exceeded(self,loop_history: list[LanguageModelMessage]) -> list[LanguageModelMessage]:
+    def _handle_token_limit_exceeded(
+        self, loop_history: list[LanguageModelMessage]
+    ) -> list[LanguageModelMessage]:
         """Handle case where token limit is exceeded by reducing sources in tool responses."""
         self._logger.warning(
             f"Length of messages is exceeds limit of {self._language_model.token_limits.token_limit_input} tokens. "
             "Reducing number of sources per tool call.",
         )
-       
-        return self._reduce_message_length_by_reducing_sources_in_tool_response(loop_history)
+
+        return self._reduce_message_length_by_reducing_sources_in_tool_response(
+            loop_history
+        )
 
     def _replace_user_message(
         self,
@@ -176,21 +190,24 @@ class LoopTokenReducer():
                                 original_user_message,
                                 "",
                             )
-                            t["text"] = rendered_user_message_string + added_to_message_by_history
+                            t["text"] = (
+                                rendered_user_message_string
+                                + added_to_message_by_history
+                            )
                         break
             elif m.content:
-                added_to_message_by_history = m.content.replace(original_user_message, "")
+                added_to_message_by_history = m.content.replace(
+                    original_user_message, ""
+                )
                 m.content = rendered_user_message_string + added_to_message_by_history
         else:
             history = history + [
                 LanguageModelUserMessage(content=rendered_user_message_string),
             ]
         return history
-    
 
     async def _get_history_from_db(
-        self,
-        remove_from_text: Callable[[str], Awaitable[str]]
+        self, remove_from_text: Callable[[str], Awaitable[str]]
     ) -> list[LanguageModelMessage]:
         """
         Get the history of the conversation. The function will retrieve a subset of the full history based on the configuration.
@@ -199,25 +216,25 @@ class LoopTokenReducer():
             list[LanguageModelMessage]: The history
         """
         full_history = get_full_history_with_contents(
-                user_message=self._user_message,
-                chat_id=self._chat_id,
-                chat_service=self._chat_service,
-                content_service=self._content_service,
-                file_content_serialization_type=(
-                    FileContentSerialization.NONE
-                    if self._has_uploaded_content_config
-                    else FileContentSerialization.FILE_NAME
-                ),
-            )
-        
-        full_history.root = await self._clean_messages(full_history.root, remove_from_text)
-        
+            user_message=self._user_message,
+            chat_id=self._chat_id,
+            chat_service=self._chat_service,
+            content_service=self._content_service,
+            file_content_serialization_type=(
+                FileContentSerialization.NONE
+                if self._has_uploaded_content_config
+                else FileContentSerialization.FILE_NAME
+            ),
+        )
+
+        full_history.root = await self._clean_messages(
+            full_history.root, remove_from_text
+        )
+
         limited_history_messages = self._limit_to_token_window(
-            full_history.root,
-            self._max_history_tokens
+            full_history.root, self._max_history_tokens
         )
-        
-        
+
         if len(limited_history_messages) == 0:
             limited_history_messages = full_history.root[-1:]
 
@@ -241,9 +258,15 @@ class LoopTokenReducer():
         return selected_messages[::-1]
 
     async def _clean_messages(
-        self, 
-        messages: list[LanguageModelMessage | LanguageModelToolMessage | LanguageModelAssistantMessage | LanguageModelSystemMessage | LanguageModelUserMessage],
-        remove_from_text: Callable[[str], Awaitable[str]]
+        self,
+        messages: list[
+            LanguageModelMessage
+            | LanguageModelToolMessage
+            | LanguageModelAssistantMessage
+            | LanguageModelSystemMessage
+            | LanguageModelUserMessage
+        ],
+        remove_from_text: Callable[[str], Awaitable[str]],
     ) -> list[LanguageModelMessage]:
         for message in messages:
             if isinstance(message.content, str):
@@ -254,13 +277,12 @@ class LoopTokenReducer():
                 )
         return messages
 
-    def _count_tokens(self,text:str) -> int:
-
+    def _count_tokens(self, text: str) -> int:
         return len(self._encoder.encode(text))
 
     def ensure_last_message_is_user_message(self, limited_history_messages):
         """
-        As the token limit can be reached in the middle of a gpt_request, 
+        As the token limit can be reached in the middle of a gpt_request,
         we move forward to the next user message,to avoid confusing messages for the LLM
         """
         idx = 0
@@ -271,45 +293,42 @@ class LoopTokenReducer():
         # FIXME: This might reduce the history by a lot if we have a lot of tool calls / references in the history. Could make sense to summarize the messages and include
         # FIXME: We should remove chunks no longer in history from handler
         return limited_history_messages[idx:]
-    
 
     def _reduce_message_length_by_reducing_sources_in_tool_response(
-      self,
-      history: list[LanguageModelMessage],
+        self,
+        history: list[LanguageModelMessage],
     ) -> list[LanguageModelMessage]:
-      """
-      Reduce the message length by removing the last source result of each tool call.
-      If there is only one source for a tool call, the tool call message is returned unchanged.
-      """
-      history_reduced: list[LanguageModelMessage] = []
-      content_chunks_reduced: list[ContentChunk] = []
-      chunk_offset = 0
-      source_offset = 0
-
-      for message in history:
-          if self._should_reduce_message(message):
-              result = self._reduce_sources_in_tool_message(
-                  message,  # type: ignore
-                  chunk_offset,
-                  source_offset,
-              )
-              content_chunks_reduced.extend(result.reduced_chunks)
-              history_reduced.append(result.message)
-              chunk_offset = result.chunk_offset
-              source_offset = result.source_offset
-          else:
-              history_reduced.append(message)
-
-      self._reference_manager.replace(chunks=content_chunks_reduced)
-      return history_reduced
-    
+        """
+        Reduce the message length by removing the last source result of each tool call.
+        If there is only one source for a tool call, the tool call message is returned unchanged.
+        """
+        history_reduced: list[LanguageModelMessage] = []
+        content_chunks_reduced: list[ContentChunk] = []
+        chunk_offset = 0
+        source_offset = 0
+
+        for message in history:
+            if self._should_reduce_message(message):
+                result = self._reduce_sources_in_tool_message(
+                    message,  # type: ignore
+                    chunk_offset,
+                    source_offset,
+                )
+                content_chunks_reduced.extend(result.reduced_chunks)
+                history_reduced.append(result.message)
+                chunk_offset = result.chunk_offset
+                source_offset = result.source_offset
+            else:
+                history_reduced.append(message)
+
+        self._reference_manager.replace(chunks=content_chunks_reduced)
+        return history_reduced
+
     def _should_reduce_message(self, message: LanguageModelMessage) -> bool:
-      """Determine if a message should have its sources reduced."""
-      return (
-          message.role == LanguageModelMessageRole.TOOL
-          and isinstance(message, LanguageModelToolMessage)
-      )
-  
+        """Determine if a message should have its sources reduced."""
+        return message.role == LanguageModelMessageRole.TOOL and isinstance(
+            message, LanguageModelToolMessage
+        )
 
     def _reduce_sources_in_tool_message(
         self,
@@ -344,8 +363,7 @@ class LoopTokenReducer():
                 chunk_offset : chunk_offset + num_sources - 1
             ]
             self._reference_manager.replace_chunks_of_tool(
-                message.tool_call_id,
-                reduced_chunks
+                message.tool_call_id, reduced_chunks
             )
 
         # Create new message with reduced sources
@@ -359,11 +377,9 @@ class LoopTokenReducer():
             message=new_message,
             reduced_chunks=content_chunks_reduced,
             chunk_offset=chunk_offset + num_sources,
-            source_offset=source_offset
-            + num_sources
-            - (1 if num_sources != 1 else 0),
+            source_offset=source_offset + num_sources - (1 if num_sources != 1 else 0),
         )
-    
+
     def _create_tool_call_message_with_reduced_sources(
         self,
         message: LanguageModelToolMessage,
@@ -384,7 +400,7 @@ class LoopTokenReducer():
         return self._create_reduced_standard_sources_message(
             message, content_chunks, source_offset
         )
-    
+
     def _create_reduced_table_search_message(
         self,
         message: LanguageModelToolMessage,
@@ -405,9 +421,7 @@ class LoopTokenReducer():
             elif isinstance(message.content, dict):
                 content_dict = message.content
             else:
-                raise ValueError(
-                    f"Unexpected content type: {type(message.content)}"
-                )
+                raise ValueError(f"Unexpected content type: {type(message.content)}")
 
             content = json.dumps(
                 {
@@ -422,7 +436,6 @@ class LoopTokenReducer():
             name=message.name,
         )
 
-
     def _create_reduced_empty_sources_message(
         self,
         message: LanguageModelToolMessage,
@@ -434,7 +447,6 @@ class LoopTokenReducer():
             name=message.name,
         )
 
-
     def _create_reduced_standard_sources_message(
         self,
         message: LanguageModelToolMessage,

unique_toolkit/history_manager/utils.py

@@ -13,7 +13,6 @@ from unique_toolkit.tools.utils.source_handling.schema import (
     SourceFormatConfig,
 )
 
-
 logger = logging.getLogger(__name__)
 
 

unique_toolkit/language_model/infos.py

@@ -875,7 +875,7 @@ class LanguageModelInfo(BaseModel):
                     deprecated_at=date(2026, 8, 7),
                     retirement_at=date(2026, 8, 7),
                 )
-            
+
             case _:
                 if isinstance(model_name, LanguageModelName):
                     raise ValueError(

unique_toolkit/language_model/schemas.py

@@ -136,7 +136,6 @@ class LanguageModelStreamResponse(BaseModel):
         """
         return not self.message.original_text and not self.tool_calls
 
-
     def to_openai_param(self) -> ChatCompletionAssistantMessageParam:
         return ChatCompletionAssistantMessageParam(
             role="assistant",

unique_toolkit/postprocessor/postprocessor_manager.py

@@ -1,5 +1,5 @@
-from abc import ABC
 import asyncio
+from abc import ABC
 from logging import Logger
 
 from unique_toolkit.chat.service import ChatService
@@ -52,8 +52,6 @@ class PostprocessorManager:
     The PostprocessorManager serves as a centralized system for managing and applying postprocessing logic to enhance response quality and consistency.
     """
 
-    
-
     def __init__(
         self,
         logger: Logger,

unique_toolkit/reference_manager/reference_manager.py

@@ -49,7 +49,6 @@ class ReferenceManager:
 
     def get_tool_chunks(self) -> dict[str, tool_chunks]:
         return self._tool_chunks
-    
 
     def get_chunks_of_all_tools(self) -> list[list[ContentChunk]]:
         return [tool_chunks.chunks for tool_chunks in self._tool_chunks.values()]
@@ -57,12 +56,12 @@ class ReferenceManager:
     def get_chunks_of_tool(self, tool_call_id: str) -> list[ContentChunk]:
         return self._tool_chunks.get(tool_call_id, tool_chunks("", [])).chunks
 
-    
-    def replace_chunks_of_tool(self, tool_call_id: str,chunks: list[ContentChunk]) -> None:
+    def replace_chunks_of_tool(
+        self, tool_call_id: str, chunks: list[ContentChunk]
+    ) -> None:
         if tool_call_id in self._tool_chunks:
             self._tool_chunks[tool_call_id].chunks = chunks
 
-
     def replace(self, chunks: list[ContentChunk]):
         self._chunks = chunks
 

unique_toolkit/short_term_memory/persistent_short_term_memory_manager.py

@@ -4,11 +4,11 @@ from logging import getLogger
 from typing import Generic, Type, TypeVar
 
 from pydantic import BaseModel
+
 from unique_toolkit.short_term_memory.schemas import ShortTermMemory
 from unique_toolkit.short_term_memory.service import ShortTermMemoryService
 from unique_toolkit.tools.utils.execution.execution import SafeTaskExecutor
 
-
 TSchema = TypeVar("TSchema", bound=BaseModel)
 
 
@@ -50,6 +50,7 @@ class PersistentShortMemoryManager(Generic[TSchema]):
 
     The PersistentShortMemoryManager is designed to handle short-term memory efficiently, ensuring data integrity and optimized storage.
     """
+
     def __init__(
         self,
         short_term_memory_service: ShortTermMemoryService,

unique_toolkit/thinking_manager/thinking_manager.py

@@ -1,9 +1,9 @@
 from logging import Logger
+
 from pydantic import BaseModel, Field
 
 from unique_toolkit.chat.service import ChatService
 from unique_toolkit.language_model.schemas import (
-    LanguageModelAssistantMessage,
     LanguageModelStreamResponse,
 )
 from unique_toolkit.tools.tool_progress_reporter import (
@@ -36,6 +36,7 @@ class ThinkingManager:
 
     The ThinkingManager enhances transparency and user understanding by providing a clear view of the assistant's reasoning process.
     """
+
     def __init__(
         self,
         logger: Logger,

unique_toolkit/tools/config.py

@@ -1,8 +1,7 @@
 from enum import StrEnum
+from typing import TYPE_CHECKING, Any
+
 import humps
-from typing import Any
-from pydantic.fields import ComputedFieldInfo, FieldInfo
-from pydantic.alias_generators import to_camel
 from pydantic import (
     BaseModel,
     ConfigDict,
@@ -10,8 +9,8 @@ from pydantic import (
     ValidationInfo,
     model_validator,
 )
-
-from typing import TYPE_CHECKING
+from pydantic.alias_generators import to_camel
+from pydantic.fields import ComputedFieldInfo, FieldInfo
 
 if TYPE_CHECKING:
     from unique_toolkit.tools.schemas import BaseToolConfig

unique_toolkit/tools/factory.py

@@ -1,18 +1,12 @@
-from typing import Callable
-
-
-from unique_toolkit.tools.tool import Tool
+from typing import TYPE_CHECKING, Callable
 
 from unique_toolkit.tools.schemas import BaseToolConfig
+from unique_toolkit.tools.tool import Tool
 
-
-from typing import TYPE_CHECKING
 if TYPE_CHECKING:
     from unique_toolkit.tools.config import ToolBuildConfig
 
 
-
-
 class ToolFactory:
     tool_map: dict[str, type[Tool]] = {}
     tool_config_map: dict[str, Callable] = {}

unique_toolkit/tools/schemas.py

@@ -4,8 +4,8 @@ import re
 from typing import Any, Optional
 
 from pydantic import BaseModel, ConfigDict, Field, field_serializer, field_validator
-from unique_toolkit.content.schemas import ContentChunk
 
+from unique_toolkit.content.schemas import ContentChunk
 from unique_toolkit.tools.config import get_configuration_dict
 from unique_toolkit.tools.utils.source_handling.schema import SourceFormatConfig
 

unique_toolkit/tools/test/test_tool_progress_reporter.py

@@ -1,6 +1,7 @@
 from unittest.mock import AsyncMock
 
 import pytest
+
 from unique_toolkit.chat.service import ChatService
 from unique_toolkit.content.schemas import ContentReference
 from unique_toolkit.language_model.schemas import LanguageModelFunction

unique_toolkit/tools/tool.py

@@ -1,24 +1,20 @@
 from abc import ABC, abstractmethod
-from enum import StrEnum
 from logging import getLogger
-from typing import Generic, TypeVar
-from typing import Any, cast
+from typing import Any, Generic, TypeVar, cast
 
-from pydantic import Field
 from typing_extensions import deprecated
+
 from unique_toolkit.app.schemas import ChatEvent
 from unique_toolkit.chat.service import (
     ChatService,
 )
+from unique_toolkit.evals.schemas import EvaluationMetricName
 from unique_toolkit.language_model import LanguageModelToolDescription
 from unique_toolkit.language_model.schemas import (
     LanguageModelFunction,
     LanguageModelMessage,
 )
 from unique_toolkit.language_model.service import LanguageModelService
-
-
-from unique_toolkit.evals.schemas import EvaluationMetricName
 from unique_toolkit.tools.config import ToolBuildConfig, ToolSelectionPolicy
 from unique_toolkit.tools.schemas import BaseToolConfig, ToolCallResponse, ToolPrompts
 from unique_toolkit.tools.tool_progress_reporter import ToolProgressReporter

unique_toolkit/tools/tool_manager.py

@@ -1,7 +1,10 @@
 import asyncio
 from logging import Logger, getLogger
+
 from pydantic import BaseModel, Field
+
 from unique_toolkit.app.schemas import ChatEvent
+from unique_toolkit.evals.schemas import EvaluationMetricName
 from unique_toolkit.language_model.schemas import (
     LanguageModelFunction,
     LanguageModelTool,
@@ -13,7 +16,6 @@ from unique_toolkit.tools.schemas import ToolCallResponse, ToolPrompts
 from unique_toolkit.tools.tool import Tool
 from unique_toolkit.tools.tool_progress_reporter import ToolProgressReporter
 from unique_toolkit.tools.utils.execution.execution import Result, SafeTaskExecutor
-from unique_toolkit.evals.schemas import EvaluationMetricName
 
 
 class ForcedToolOption:
@@ -36,7 +38,6 @@ class ToolManagerConfig(BaseModel):
     )
 
 
-
 class ToolManager:
     """
     Manages the tools available to the agent and executes tool calls.

unique_toolkit/tools/tool_progress_reporter.py

@@ -5,6 +5,7 @@ from functools import wraps
 from typing import Protocol
 
 from pydantic import BaseModel
+
 from unique_toolkit.chat.service import ChatService
 from unique_toolkit.content.schemas import ContentReference
 from unique_toolkit.language_model.schemas import (

unique_toolkit/tools/utils/source_handling/schema.py

@@ -3,7 +3,6 @@ from pydantic import BaseModel
 
 from unique_toolkit.tools.config import get_configuration_dict
 
-
 SOURCE_TEMPLATE = "<source${index}>${document}${info}${text}</source${index}>"
 SECTIONS = {
     "document": "<|document|>{}<|/document|>\n",

unique_toolkit/tools/utils/source_handling/tests/test_source_formatting.py

@@ -1,4 +1,5 @@
 import pytest
+
 from unique_toolkit.content.schemas import ContentChunk
 from unique_toolkit.tools.utils.source_handling.schema import SourceFormatConfig
 from unique_toolkit.tools.utils.source_handling.source_formatting import (