unique_toolkit 1.62.3__tar.gz → 1.63.0__tar.gz

{unique_toolkit-1.62.3 → unique_toolkit-1.63.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: unique-toolkit
-Version: 1.62.3
+Version: 1.63.0
 Summary: 
 Author: Cedric Klinkert, Martin Fadler, Sadique Sheik, Fabian Schläpfer, Pascal Hauri, Ejiroghene Oki-Peter, Rami Azouz, Ricardo Visini, Sanzio Monti
 Author-email: Cedric Klinkert <cedric.klinkert@unique.ch>, Martin Fadler <martin.fadler@unique.ch>, Sadique Sheik <sadique@unique.ch>, Fabian Schläpfer <fabian@unique.ch>, Pascal Hauri <pascal@unique.ch>, Ejiroghene Oki-Peter <oki-peter.ejiroghene@unique.ai>, Rami Azouz <rami.ext@unique.ch>, Ricardo Visini <ricardo.visini@unique.ch>, Sanzio Monti <sanzio.monti@unique.ai>

{unique_toolkit-1.62.3 → unique_toolkit-1.63.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "unique_toolkit"
-version = "1.62.3"
+version = "1.63.0"
 description = ""
 readme = "README.md"
 requires-python = ">=3.12"

{unique_toolkit-1.62.3 → unique_toolkit-1.63.0}/unique_toolkit/agentic/history_manager/history_construction_with_contents.py

@@ -9,6 +9,7 @@ from itertools import groupby
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
+    from unique_toolkit.content.schemas import ContentChunk
     from unique_toolkit.language_model.builder import MessagesBuilder
 
 import numpy as np
@@ -256,13 +257,30 @@ def get_full_history_with_contents_and_tool_calls(
     content_service: ContentService,
     include_images: ImageContentInclusion = ImageContentInclusion.ALL,
     file_content_serialization_type: FileContentSerialization = FileContentSerialization.FILE_NAME,
-) -> LanguageModelMessages:
-    """Build the full LLM message history, including persisted tool call rounds."""
+) -> tuple[LanguageModelMessages, int, dict[int, "ContentChunk"]]:
+    """Build the full LLM message history, including persisted tool call rounds.
+
+    Returns:
+        A triple of (messages, max_source_number, source_map) where
+        *max_source_number* is the highest ``source_number`` found in any
+        persisted tool response (``-1`` when none exist) and *source_map*
+        maps each ``source_number`` to a ``ContentChunk`` reconstructed from
+        the persisted response content.
+    """
+    from unique_toolkit.agentic.history_manager.utils import (
+        build_source_map_from_tool_calls,
+        compute_max_source_number_from_tool_calls,
+    )
+    from unique_toolkit.content.schemas import (
+        ContentChunk as ContentChunk,  # noqa: F811
+    )
+
     chat_history = chat_service.get_full_history()
 
     assistant_message_ids = [
         msg.id for msg in chat_history if msg.role == ChatRole.ASSISTANT and msg.id
     ]
+    all_tool_calls: list[ChatMessageTool] = []
     tool_calls_by_message: dict[str, list[ChatMessageTool]] = {}
     if assistant_message_ids:
         try:
@@ -276,6 +294,10 @@ def get_full_history_with_contents_and_tool_calls(
             _LOGGER.warning(
                 "Failed to batch-load tool calls, falling back to empty", exc_info=True
             )
+    max_source_number = compute_max_source_number_from_tool_calls(all_tool_calls)
+    source_map: dict[int, ContentChunk] = build_source_map_from_tool_calls(
+        all_tool_calls
+    )
 
     grouped_elements = get_chat_history_with_contents(
         user_message=user_message,
@@ -366,7 +388,7 @@ def get_full_history_with_contents_and_tool_calls(
             content_service=content_service,
             chat_id=chat_id,
         )
-    return builder.build()
+    return builder.build(), max_source_number, source_map
 
 
 def get_full_history_as_llm_messages(

{unique_toolkit-1.62.3 → unique_toolkit-1.63.0}/unique_toolkit/agentic/history_manager/history_manager.py

@@ -10,12 +10,16 @@ from unique_toolkit._common.feature_flags.schema import (
 )
 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.history_manager.utils import (
+    serialize_tool_content_json,
+    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.chat.schemas import ChatMessageTool, ChatMessageToolResponse
+from unique_toolkit.content.schemas import ContentChunk
 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 (
@@ -64,13 +68,6 @@ class HistoryManagerConfig(BaseModel):
         description="The fraction of the max input tokens that will be reserved for the history.",
     )
 
-    percent_for_tool_call_history: float = Field(
-        default=0.0,
-        ge=0.0,
-        lt=1.0,
-        description="The fraction of the max input tokens reserved for tool call rounds. 0 disables.",
-    )
-
     language_model: LMI = LanguageModelInfo.from_name(DEFAULT_GPT_4o)
 
     @property
@@ -80,13 +77,6 @@ class HistoryManagerConfig(BaseModel):
             * self.percent_of_max_tokens_for_history,
         )
 
-    @property
-    def max_tool_call_history_tokens(self) -> int:
-        return int(
-            self.language_model.token_limits.token_limit_input
-            * self.percent_for_tool_call_history,
-        )
-
     uploaded_content_config: (
         Annotated[
             UploadedContentConfig,
@@ -135,12 +125,15 @@ class HistoryManager:
             has_uploaded_content_config=bool(self._config.uploaded_content_config),
             language_model=self._language_model,
             reference_manager=reference_manager,
-            max_tokens_for_tool_call_history=self._config.max_tool_call_history_tokens,
         )
+        self._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
+        self._initial_source_offset = 0
+        self._db_source_map: dict[int, ContentChunk] = {}
+        self._source_offset_initialized = False
         self._collected_tool_response_image_urls: list[tuple[str, str]] = []
 
     def add_tool_call(self, tool_call: LanguageModelFunction) -> None:
@@ -193,16 +186,13 @@ class HistoryManager:
                 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,
             )
             content = stringified_sources
 
-            self._source_enumerator += len(
-                sources
-            )  # To make sure all sources have unique source numbers
+            self._source_enumerator += len(sources)
 
         if tool_response.system_reminder:
             content += f"\n\n{tool_response.system_reminder}"
@@ -243,6 +233,18 @@ class HistoryManager:
             image_data_urls_from_tools=image_data_from_tools,
         )
         self._collected_tool_response_image_urls = []
+
+        if not self._source_offset_initialized:
+            offset = max(0, self._token_reducer.max_db_source_number + 1)
+            self._source_enumerator = offset
+            self._initial_source_offset = offset
+            self._db_source_map = self._token_reducer.db_source_map
+            self._source_offset_initialized = True
+
+        self._source_enumerator = self._initial_source_offset + len(
+            self._reference_manager.get_chunks()
+        )
+
         return messages
 
     async def get_user_visible_chat_history(
@@ -267,29 +269,47 @@ class HistoryManager:
             )
         return LanguageModelMessages(history)
 
+    @staticmethod
+    def _placeholder_chunk() -> ContentChunk:
+        """An empty but backend-valid placeholder for gap positions."""
+        return ContentChunk(key="", chunk_id="")
+
+    def get_content_chunks_for_backend(self) -> list[ContentChunk]:
+        """Build a positional list where ``result[N]`` is the chunk for ``[sourceN]``.
+
+        Positions ``0..K-1`` are filled from ``_db_source_map`` (prior turns).
+        Gaps get placeholder ``ContentChunk`` instances with empty-string
+        ``key`` / ``chunk_id`` so the backend doesn't reject the payload.
+        Positions ``K..`` are the current turn's chunks from the reference
+        manager.
+        """
+        current_chunks = self._reference_manager.get_chunks()
+        total_size = self._initial_source_offset + len(current_chunks)
+        result: list[ContentChunk] = [
+            self._placeholder_chunk() for _ in range(total_size)
+        ]
+        for source_number, chunk in self._db_source_map.items():
+            if source_number < self._initial_source_offset:
+                result[source_number] = chunk
+        for i, chunk in enumerate(current_chunks):
+            result[self._initial_source_offset + i] = chunk
+        return result
+
     def extract_message_tools(self) -> list[ChatMessageTool]:
         """Convert the in-memory loop history into persistable ChatMessageTool records."""
+        # Build a map of tool_call_id -> response_content in O(T) time
+        tool_responses: dict[str, str | None] = {}
+        for msg in self._loop_history:
+            if isinstance(msg, LanguageModelToolMessage):
+                response_content = msg.content if isinstance(msg.content, str) else None
+                tool_responses[msg.tool_call_id] = response_content
+
         records: list[ChatMessageTool] = []
         round_index = 0
-        i = 0
-        while i < len(self._loop_history):
-            msg = self._loop_history[i]
+        for msg in self._loop_history:
             if isinstance(msg, LanguageModelAssistantMessage) and msg.tool_calls:
                 for seq_index, tc in enumerate(msg.tool_calls):
-                    response_content = None
-                    for j in range(i + 1, len(self._loop_history)):
-                        candidate = self._loop_history[j]
-                        if (
-                            isinstance(candidate, LanguageModelToolMessage)
-                            and candidate.tool_call_id == tc.id
-                        ):
-                            response_content = (
-                                candidate.content
-                                if isinstance(candidate.content, str)
-                                else None
-                            )
-                            break
-
+                    response_content = tool_responses.get(tc.id) if tc.id else None
                     response = (
                         ChatMessageToolResponse(content=response_content)
                         if response_content is not None
@@ -306,7 +326,6 @@ class HistoryManager:
                         )
                     )
                 round_index += 1
-            i += 1
         return records
 
     @staticmethod
@@ -366,4 +385,4 @@ def _strip_uncited_sources_from_content(content: str, cited: set[int]) -> str:
         return content
 
     filtered = [item for item in data if item.get("source_number") in cited]
-    return json.dumps(filtered)
+    return serialize_tool_content_json(filtered)

{unique_toolkit-1.62.3 → unique_toolkit-1.63.0}/unique_toolkit/agentic/history_manager/loop_token_reducer.py

@@ -12,6 +12,7 @@ from unique_toolkit.agentic.history_manager.history_construction_with_contents i
     FileContentSerialization,
     get_full_history_with_contents_and_tool_calls,
 )
+from unique_toolkit.agentic.history_manager.utils import serialize_tool_content_json
 from unique_toolkit.agentic.reference_manager.reference_manager import ReferenceManager
 from unique_toolkit.app.schemas import ChatEvent
 from unique_toolkit.app.unique_settings import UniqueSettings
@@ -53,11 +54,9 @@ class LoopTokenReducer:
         has_uploaded_content_config: bool,
         reference_manager: ReferenceManager,
         language_model: LMI,
-        max_tokens_for_tool_call_history: int | None = None,
     ):
         self._max_history_tokens = max_history_tokens
         self._has_uploaded_content_config = has_uploaded_content_config
-        self._max_tokens_for_tool_call_history = max_tokens_for_tool_call_history
         self._logger = logger
         self._reference_manager = reference_manager
         self._language_model = language_model
@@ -72,6 +71,16 @@ class LoopTokenReducer:
             self._language_model.token_limits.token_limit_input
             * (1 - MAX_INPUT_TOKENS_SAFETY_PERCENTAGE)
         )
+        self._max_db_source_number: int = -1
+        self._db_source_map: dict[int, ContentChunk] = {}
+
+    @property
+    def max_db_source_number(self) -> int:
+        return self._max_db_source_number
+
+    @property
+    def db_source_map(self) -> dict[int, ContentChunk]:
+        return self._db_source_map
 
     def _get_encoder(self, language_model: LMI) -> Callable[[str], list[int]]:
         return language_model.get_encoder()
@@ -105,7 +114,9 @@ class LoopTokenReducer:
 
         while self._exceeds_token_limit(token_count):
             token_count_before_reduction = token_count
-            loop_history = self._handle_token_limit_exceeded(loop_history, token_count)
+            loop_history[:] = self._handle_token_limit_exceeded(
+                loop_history, token_count
+            )
             messages = self._construct_history(
                 history_from_db,
                 loop_history,
@@ -292,7 +303,7 @@ class LoopTokenReducer:
         Returns:
             list[LanguageModelMessage]: The history
         """
-        full_history = get_full_history_with_contents_and_tool_calls(
+        full_history, max_src, src_map = get_full_history_with_contents_and_tool_calls(
             user_message=self._user_message,
             chat_id=self._chat_id,
             chat_service=self._chat_service,
@@ -303,16 +314,14 @@ class LoopTokenReducer:
                 else FileContentSerialization.FILE_NAME
             ),
         )
+        self._max_db_source_number = max_src
+        self._db_source_map = src_map
+
         if remove_from_text is not None:
             full_history.root = await self._clean_messages(
                 full_history.root, remove_from_text
             )
 
-        if self._max_tokens_for_tool_call_history is not None:
-            full_history.root = self._limit_tool_call_tokens(
-                full_history.root, self._max_tokens_for_tool_call_history
-            )
-
         limited_history_messages = self._limit_to_token_window(
             full_history.root, self._max_history_tokens
         )
@@ -383,77 +392,6 @@ class LoopTokenReducer:
 
         return [msg for turn in selected_turns[::-1] for msg in turn]
 
-    def _limit_tool_call_tokens(
-        self,
-        messages: list[LanguageModelMessage],
-        budget: int,
-    ) -> list[LanguageModelMessage]:
-        """Drop older tool-call rounds so the total stays within *budget* tokens.
-
-        A "round" is a contiguous block of one ``LanguageModelAssistantMessage``
-        (with ``tool_calls``) followed by its ``LanguageModelToolMessage`` replies.
-        Rounds are kept from the back (most recent first).  When *budget* is
-        positive, at least one round is always preserved even if it alone exceeds
-        the budget.  A *budget* of ``0`` removes **all** tool-call rounds.
-        """
-        rounds: list[tuple[int, int]] = []
-        i = 0
-        while i < len(messages):
-            msg = messages[i]
-            if isinstance(msg, LanguageModelAssistantMessage) and msg.tool_calls:
-                start = i
-                i += 1
-                while i < len(messages) and isinstance(
-                    messages[i], LanguageModelToolMessage
-                ):
-                    i += 1
-                rounds.append((start, i))
-            else:
-                i += 1
-
-        if not rounds:
-            return messages
-
-        if budget <= 0:
-            drop_all: set[int] = set()
-            for start, end in rounds:
-                drop_all.update(range(start, end))
-            self._logger.info(
-                f"Tool-call budget is 0: dropping all {len(rounds)} tool-call rounds",
-            )
-            return [m for i, m in enumerate(messages) if i not in drop_all]
-
-        round_tokens = [
-            sum(
-                self._count_message_tokens(LanguageModelMessages(root=[m]))
-                for m in messages[start:end]
-            )
-            for start, end in rounds
-        ]
-
-        keep: set[int] = set()
-        cumulative = 0
-        for idx in range(len(rounds) - 1, -1, -1):
-            cumulative += round_tokens[idx]
-            if cumulative > budget and keep:
-                break
-            keep.add(idx)
-
-        drop_indices: set[int] = set()
-        for idx, (start, end) in enumerate(rounds):
-            if idx not in keep:
-                drop_indices.update(range(start, end))
-
-        if not drop_indices:
-            return messages
-
-        self._logger.info(
-            f"Tool-call budget: dropping {len(rounds) - len(keep)} of {len(rounds)} "
-            f"rounds ({sum(round_tokens[i] for i in range(len(rounds)) if i not in keep)} tokens freed)",
-        )
-
-        return [m for i, m in enumerate(messages) if i not in drop_indices]
-
     async def _clean_messages(
         self,
         messages: list[
@@ -515,7 +453,7 @@ class LoopTokenReducer:
         history_reduced: list[LanguageModelMessage] = []
         content_chunks_reduced: list[ContentChunk] = []
         chunk_offset = 0
-        source_offset = 0
+        source_offset = max(0, self._max_db_source_number + 1)
 
         for message in history:
             if self._should_reduce_message(message):
@@ -651,7 +589,7 @@ class LoopTokenReducer:
             else:
                 raise ValueError(f"Unexpected content type: {type(message.content)}")
 
-            content = json.dumps(
+            content = serialize_tool_content_json(
                 {
                     "source_number": source_offset,
                     "content": content_dict.get("content"),
@@ -692,7 +630,7 @@ class LoopTokenReducer:
         ]
 
         return LanguageModelToolMessage(
-            content=json.dumps(sources),
+            content=serialize_tool_content_json(sources),
             tool_call_id=message.tool_call_id,
             name=message.name,
         )

{unique_toolkit-1.62.3 → unique_toolkit-1.63.0}/unique_toolkit/agentic/history_manager/utils.py

@@ -14,6 +14,11 @@ from unique_toolkit.language_model.schemas import (
 logger = logging.getLogger(__name__)
 
 
+def serialize_tool_content_json(value: Any) -> str:
+    """Serialize tool-role JSON text while preserving readable Unicode."""
+    return json.dumps(value, ensure_ascii=False)
+
+
 def convert_tool_interactions_to_content_messages(
     loop_history: list[LanguageModelMessage],
 ) -> list[LanguageModelMessage]:
@@ -80,7 +85,7 @@ def transform_chunks_to_string(
         }
         for i, chunk in enumerate(content_chunks)
     ]
-    return json.dumps(sources), sources
+    return serialize_tool_content_json(sources), sources
 
 
 def load_sources_from_string(
@@ -95,3 +100,70 @@ def load_sources_from_string(
     except (json.JSONDecodeError, ValueError):
         logger.warning("Failed to parse source string")
         return None
+
+
+def _parse_sources_from_response(content: str) -> list[dict[str, Any]]:
+    """Parse the JSON source array from a persisted tool response content string.
+
+    Returns an empty list when the content is not valid JSON or not a list.
+    """
+    try:
+        data = json.loads(content)
+        if isinstance(data, list):
+            return data
+        return []
+    except (json.JSONDecodeError, ValueError) as exc:
+        logger.warning(
+            "_parse_sources_from_response: JSON parse failed (content_len=%d): %s",
+            len(content),
+            exc,
+        )
+        return []
+
+
+def compute_max_source_number_from_tool_calls(
+    tool_calls: list[Any],
+) -> int:
+    """Find the highest ``source_number`` across all persisted tool-call responses.
+
+    Returns ``-1`` when no source numbers are found.
+    """
+    from unique_toolkit.chat.schemas import ChatMessageTool
+
+    max_num = -1
+    for tc in tool_calls:
+        if not isinstance(tc, ChatMessageTool):
+            continue
+        if not tc.response or not tc.response.content:
+            continue
+        for entry in _parse_sources_from_response(tc.response.content):
+            sn = entry.get("source_number")
+            if isinstance(sn, int) and sn > max_num:
+                max_num = sn
+    return max_num
+
+
+def build_source_map_from_tool_calls(
+    tool_calls: list[Any],
+) -> dict[int, ContentChunk]:
+    """Build a mapping ``{source_number: ContentChunk}`` from persisted tool responses.
+
+    Only entries that have both a valid ``source_number`` and ``content`` are included.
+    """
+    from unique_toolkit.chat.schemas import ChatMessageTool
+
+    source_map: dict[int, ContentChunk] = {}
+    for tc in tool_calls:
+        if not isinstance(tc, ChatMessageTool):
+            continue
+        if not tc.response or not tc.response.content:
+            continue
+        for entry in _parse_sources_from_response(tc.response.content):
+            sn = entry.get("source_number")
+            text = entry.get("content")
+            if isinstance(sn, int) and isinstance(text, str):
+                content_id = entry.get("content_id") or ""
+                source_map[sn] = ContentChunk(
+                    id=content_id, text=text, key="", chunk_id=""
+                )
+    return source_map

{unique_toolkit-1.62.3 → unique_toolkit-1.63.0}/unique_toolkit/chat/functions.py

@@ -1410,9 +1410,6 @@ async def update_message_execution_async(
         raise e
 
 
-_MAX_TOOL_RESPONSE_CONTENT_LEN = 10_000
-
-
 def _build_tool_call_items(
     tool_calls: list[ChatMessageTool],
 ) -> list[dict[str, object]]:
@@ -1423,15 +1420,7 @@ def _build_tool_call_items(
             "arguments": tc.arguments,
             "roundIndex": tc.round_index,
             "sequenceIndex": tc.sequence_index,
-            **(
-                {
-                    "response": {
-                        "content": tc.response.content[:_MAX_TOOL_RESPONSE_CONTENT_LEN]
-                    }
-                }
-                if tc.response
-                else {}
-            ),
+            **({"response": {"content": tc.response.content}} if tc.response else {}),
         }
         for tc in tool_calls
     ]
@@ -1484,7 +1473,8 @@ async def create_message_tools_async(
             messageId=message_id,
             tools=_build_tool_call_items(tool_calls),  # type: ignore
         )
-        return [ChatMessageTool.model_validate(dict(item)) for item in result.data]
+        tools = [ChatMessageTool.model_validate(dict(item)) for item in result.data]
+        return tools
     except Exception as e:
         _LOGGER.error(f"Failed to create message tools: {e}")
         raise e
@@ -1507,7 +1497,8 @@ def get_message_tools(
             company_id=company_id,
             messageIds=ids,
         )
-        return [ChatMessageTool.model_validate(dict(item)) for item in result.data]
+        tools = [ChatMessageTool.model_validate(dict(item)) for item in result.data]
+        return tools
     except Exception as e:
         _LOGGER.error(f"Failed to get message tools: {e}")
         raise e