Repository: microsoft/autogen
Pull request: Min tokens in token limiter
## Why are these changes needed?

<!-- Please give a short summary of the change and the problem this solves. -->
A `min_tokens_threshold` is added to `MessageTokenLimiter`, which gives the option to cut tokens only after the specified message limit is reached.

## Related issue number

Closes #2306.

## Checks

- [x] I've included any doc changes needed for https://microsoft.github.io/autogen/. See https://microsoft.github.io/autogen/docs/Contribute#documentation to build and test documentation locally.
- [x] I've added tests (if relevant) corresponding to the changes introduced in this PR.
- [x] I've made sure all auto checks have passed.


Problem statement:
[Feature Request]: MessageTransform when token limit is exceeded.
### Is your feature request related to a problem? Please describe.

You can either execute a transformation from the beginning of a conversation or not at all. I think it would be useful to execute a transformation (transformations at this moment refer to content reduction), when the token limit is (or is close to be) exceeded.

### Describe the solution you'd like

Give the option to call transformations when the token limit is (or is close to) exceeded.

### Additional context

Let me know if you want me to open a PR.

Hints from issue discussion:
You should be able to create a custom transformation to get the specific behavior.

```python
from autogen.agentchat.contrib.capabilities import transform_messages, transforms
from typing import Dict, List

class CustomTruncation:
    def __init__(self, ...):
        self._token_limiter = transform.MessageTokenLimiter(...)
    ...

    def apply_transform(self, messages: List[Dict]) -> List[Dict]:
        if close_to_limit():
            self._token_limiter.apply_transform(messages)
        else:
            return messages
```

> You should be able to create a custom transformation to get the specific behavior.
> 
> ```python
> from autogen.agentchat.contrib.capabilities import transform_messages, transforms
> from typing import Dict, List
> 
> class CustomTruncation:
>     def __init__(self, ...):
>         self._token_limiter = transform.MessageTokenLimiter(...)
>     ...
> 
>     def apply_transform(self, messages: List[Dict]) -> List[Dict]:
>         if close_to_limit():
>             self._token_limiter.apply_transform(messages)
>         else:
>             return messages
> ```

Sure, I just think it should be easier than that.
I have two solutions in mind:
1. Add "min_token_threshold: Optional[int] = None" to the constructor: Would apply the transform if the threshold is exceeded. Simple to implement.
2. Add "min_token_threshold: Optional[Callable] = None" to the constructor: The callable returns a boolean to decide if transform should be applied. A lot more flexible, the user can dictate the behavior. Not sure if needed


> I have two solutions in mind:
> 
> 1. Add "min_token_threshold: Optional[int] = None" to the constructor: Would apply the transform if the threshold is exceeded. Simple to implement.

I don't like it very much, because the user would have to find the token limit, or call the corresponding function to get it, which could be abstracted.

> 2. Add "min_token_threshold: Optional[Callable] = None" to the constructor: The callable returns a boolean to decide if transform should be applied. A lot more flexible, the user can dictate the behavior. Not sure if needed

I like this option very much, and this is an approach that could be extended further than the minimum token case. My additional thought is that probably it could by a `when_to_execute: Optional[Union[Callable, str]] = None` argument, that could also offer some pre-implemented cases (like the token threshold). This way the users don't have to reimplement some standard functions.
Sounds good, I have a couple of PRs open, and once I merge those in I'll add it to the list (more than happy if you want to take over this PR)
> Sounds good, I have a couple of PRs open, and once I merge those in I'll add it to the list (more than happy if you want to take over this PR)

I would love to open this PR!

Commit message: Update docs and notebook

Patch to review:
diff --git a/autogen/agentchat/contrib/capabilities/transforms.py b/autogen/agentchat/contrib/capabilities/transforms.py
index 6dc1d59fe9c7..38474cc13efc 100644
--- a/autogen/agentchat/contrib/capabilities/transforms.py
+++ b/autogen/agentchat/contrib/capabilities/transforms.py
@@ -51,8 +51,7 @@ class MessageHistoryLimiter:
     def __init__(self, max_messages: Optional[int] = None):
         """
         Args:
-            max_messages (None or int): Maximum number of messages to keep in the context.
-            Must be greater than 0 if not None.
+            max_messages Optional[int]: Maximum number of messages to keep in the context. Must be greater than 0 if not None.
         """
         self._validate_max_messages(max_messages)
         self._max_messages = max_messages
@@ -70,6 +69,7 @@ def apply_transform(self, messages: List[Dict]) -> List[Dict]:
         Returns:
             List[Dict]: A new list containing the most recent messages up to the specified maximum.
         """
+
         if self._max_messages is None:
             return messages
 
@@ -108,13 +108,15 @@ class MessageTokenLimiter:
 
     The truncation process follows these steps in order:
 
-    1. Messages are processed in reverse order (newest to oldest).
-    2. Individual messages are truncated based on max_tokens_per_message. For multimodal messages containing both text
+    1. The minimum tokens threshold (`min_tokens`) is checked (0 by default). If the total number of tokens in messages
+        are less than this threshold, then the messages are returned as is. In other case, the following process is applied.
+    2. Messages are processed in reverse order (newest to oldest).
+    3. Individual messages are truncated based on max_tokens_per_message. For multimodal messages containing both text
         and other types of content, only the text content is truncated.
-    3. The overall conversation history is truncated based on the max_tokens limit. Once the accumulated token count
+    4. The overall conversation history is truncated based on the max_tokens limit. Once the accumulated token count
         exceeds this limit, the current message being processed get truncated to meet the total token count and any
         remaining messages get discarded.
-    4. The truncated conversation history is reconstructed by prepending the messages to a new list to preserve the
+    5. The truncated conversation history is reconstructed by prepending the messages to a new list to preserve the
         original message order.
     """
 
@@ -122,6 +124,7 @@ def __init__(
         self,
         max_tokens_per_message: Optional[int] = None,
         max_tokens: Optional[int] = None,
+        min_tokens: Optional[int] = None,
         model: str = "gpt-3.5-turbo-0613",
     ):
         """
@@ -130,11 +133,14 @@ def __init__(
                 Must be greater than or equal to 0 if not None.
             max_tokens (Optional[int]): Maximum number of tokens to keep in the chat history.
                 Must be greater than or equal to 0 if not None.
+            min_tokens (Optional[int]): Minimum number of tokens in messages to apply the transformation.
+                Must be greater than or equal to 0 if not None.
             model (str): The target OpenAI model for tokenization alignment.
         """
         self._model = model
         self._max_tokens_per_message = self._validate_max_tokens(max_tokens_per_message)
         self._max_tokens = self._validate_max_tokens(max_tokens)
+        self._min_tokens = self._validate_min_tokens(min_tokens, max_tokens)
 
     def apply_transform(self, messages: List[Dict]) -> List[Dict]:
         """Applies token truncation to the conversation history.
@@ -147,6 +153,11 @@ def apply_transform(self, messages: List[Dict]) -> List[Dict]:
         """
         assert self._max_tokens_per_message is not None
         assert self._max_tokens is not None
+        assert self._min_tokens is not None
+
+        # if the total number of tokens in the messages is less than the min_tokens, return the messages as is
+        if not self._check_tokens_threshold(messages):
+            return messages
 
         temp_messages = copy.deepcopy(messages)
         processed_messages = []
@@ -194,6 +205,19 @@ def get_logs(self, pre_transform_messages: List[Dict], post_transform_messages:
             return logs_str, True
         return "No tokens were truncated.", False
 
+    def _check_tokens_threshold(self, messages: List[Dict]) -> bool:
+        """
+        Returns True if no minimum tokens restrictions are applied.
+
+        Either if the total number of tokens in the messages is greater than or equal to the `min_theshold_tokens`,
+        or no minimum tokens threshold is set.
+        """
+        if not self._min_tokens:
+            return True
+
+        messages_tokens = sum(_count_tokens(msg["content"]) for msg in messages if "content" in msg)
+        return messages_tokens >= self._min_tokens
+
     def _truncate_str_to_tokens(self, contents: Union[str, List], n_tokens: int) -> Union[str, List]:
         if isinstance(contents, str):
             return self._truncate_tokens(contents, n_tokens)
@@ -244,6 +268,15 @@ def _validate_max_tokens(self, max_tokens: Optional[int] = None) -> Optional[int
 
         return max_tokens if max_tokens is not None else sys.maxsize
 
+    def _validate_min_tokens(self, min_tokens: int, max_tokens: int) -> int:
+        if min_tokens is None:
+            return 0
+        if min_tokens < 0:
+            raise ValueError("min_tokens must be None or greater than or equal to 0")
+        if max_tokens is not None and min_tokens >= max_tokens:
+            raise ValueError("min_tokens must be less than max_tokens")
+        return min_tokens
+
 
 def _count_tokens(content: Union[str, List[Dict[str, Any]]]) -> int:
     token_count = 0
diff --git a/notebook/agentchat_transform_messages.ipynb b/notebook/agentchat_transform_messages.ipynb
index ab8bc762fc76..82d0893587c5 100644
--- a/notebook/agentchat_transform_messages.ipynb
+++ b/notebook/agentchat_transform_messages.ipynb
@@ -24,10 +24,22 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 1,
    "id": "47773f79-c0fd-4993-bc6e-3d1a57690118",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "ename": "ModuleNotFoundError",
+     "evalue": "No module named 'autogen'",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[1;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
+      "Cell \u001b[1;32mIn[1], line 7\u001b[0m\n\u001b[0;32m      4\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mre\u001b[39;00m\n\u001b[0;32m      5\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtyping\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Dict, List\n\u001b[1;32m----> 7\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mautogen\u001b[39;00m\n\u001b[0;32m      8\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mautogen\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01magentchat\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcontrib\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcapabilities\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m transform_messages, transforms\n",
+      "\u001b[1;31mModuleNotFoundError\u001b[0m: No module named 'autogen'"
+     ]
+    }
+   ],
    "source": [
     "import copy\n",
     "import os\n",
@@ -41,7 +53,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "id": "9f09246b-a7d0-4238-b62c-1e72c7d815b3",
    "metadata": {},
    "outputs": [],
@@ -95,7 +107,7 @@
     "Imagine a scenario where the LLM generates an extensive amount of text, surpassing the token limit imposed by your API provider. To address this issue, you can leverage `TransformMessages` along with its constituent transformations, `MessageHistoryLimiter` and `MessageTokenLimiter`.\n",
     "\n",
     "- `MessageHistoryLimiter`: You can restrict the total number of messages considered as context history. This transform is particularly useful when you want to limit the conversational context to a specific number of recent messages, ensuring efficient processing and response generation.\n",
-    "- `MessageTokenLimiter`: Enables you to cap the total number of tokens, either on a per-message basis or across the entire context history (or both). This transformation is invaluable when you need to adhere to strict token limits imposed by your API provider, preventing unnecessary costs or errors caused by exceeding the allowed token count."
+    "- `MessageTokenLimiter`: Enables you to cap the total number of tokens, either on a per-message basis or across the entire context history (or both). This transformation is invaluable when you need to adhere to strict token limits imposed by your API provider, preventing unnecessary costs or errors caused by exceeding the allowed token count. Additionally, a min tokens threshold can be applied, ensuring that the transformation is only applied when the number of tokens is not less than the specified threshold."
    ]
   },
   {
@@ -109,7 +121,7 @@
     "max_msg_transfrom = transforms.MessageHistoryLimiter(max_messages=3)\n",
     "\n",
     "# Limit the token limit per message to 10 tokens\n",
-    "token_limit_transform = transforms.MessageTokenLimiter(max_tokens_per_message=3)"
+    "token_limit_transform = transforms.MessageTokenLimiter(max_tokens_per_message=3, min_tokens=10)"
    ]
   },
   {
@@ -306,7 +318,7 @@
     "context_handling = transform_messages.TransformMessages(\n",
     "    transforms=[\n",
     "        transforms.MessageHistoryLimiter(max_messages=10),\n",
-    "        transforms.MessageTokenLimiter(max_tokens=1000, max_tokens_per_message=50),\n",
+    "        transforms.MessageTokenLimiter(max_tokens=1000, max_tokens_per_message=50, min_tokens=500),\n",
     "    ]\n",
     ")\n",
     "\n",
@@ -494,7 +506,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.8"
+   "version": "3.11.5"
   }
  },
  "nbformat": 4,
diff --git a/test/agentchat/contrib/capabilities/test_transforms.py b/test/agentchat/contrib/capabilities/test_transforms.py
index 1a929e4c6ba1..6d9441d53e6a 100644
--- a/test/agentchat/contrib/capabilities/test_transforms.py
+++ b/test/agentchat/contrib/capabilities/test_transforms.py
@@ -20,7 +20,7 @@ def get_short_messages() -> List[Dict]:
     return [
         {"role": "user", "content": "hello"},
         {"role": "assistant", "content": [{"type": "text", "text": "there"}]},
-        {"role": "user", "content": "how"},
+        {"role": "user", "content": "how are you"},
     ]
 
 
@@ -38,7 +38,12 @@ def message_token_limiter() -> MessageTokenLimiter:
     return MessageTokenLimiter(max_tokens_per_message=3)
 
 
-# MessageHistoryLimiter tests
+@pytest.fixture
+def message_token_limiter_with_threshold() -> MessageTokenLimiter:
+    return MessageTokenLimiter(max_tokens_per_message=1, min_tokens=10)
+
+
+# MessageHistoryLimiter
 
 
 @pytest.mark.parametrize(
@@ -71,7 +76,7 @@ def test_message_history_limiter_get_logs(message_history_limiter, messages, exp
 
 @pytest.mark.parametrize(
     "messages, expected_token_count, expected_messages_len",
-    [(get_long_messages(), 9, 5), (get_short_messages(), 3, 3), (get_no_content_messages(), 0, 2)],
+    [(get_long_messages(), 9, 5), (get_short_messages(), 5, 3), (get_no_content_messages(), 0, 2)],
 )
 def test_message_token_limiter_apply_transform(
     message_token_limiter, messages, expected_token_count, expected_messages_len
@@ -83,6 +88,20 @@ def test_message_token_limiter_apply_transform(
     assert len(transformed_messages) == expected_messages_len
 
 
+@pytest.mark.parametrize(
+    "messages, expected_token_count, expected_messages_len",
+    [(get_long_messages(), 5, 5), (get_short_messages(), 5, 3), (get_no_content_messages(), 0, 2)],
+)
+def test_message_token_limiter_with_threshold_apply_transform(
+    message_token_limiter_with_threshold, messages, expected_token_count, expected_messages_len
+):
+    transformed_messages = message_token_limiter_with_threshold.apply_transform(messages)
+    assert (
+        sum(_count_tokens(msg["content"]) for msg in transformed_messages if "content" in msg) == expected_token_count
+    )
+    assert len(transformed_messages) == expected_messages_len
+
+
 @pytest.mark.parametrize(
     "messages, expected_logs, expected_effect",
     [
@@ -102,21 +121,87 @@ def test_message_token_limiter_get_logs(message_token_limiter, messages, expecte
 if __name__ == "__main__":
     long_messages = get_long_messages()
     short_messages = get_short_messages()
+    no_content_messages = get_no_content_messages()
     message_history_limiter = MessageHistoryLimiter(max_messages=3)
     message_token_limiter = MessageTokenLimiter(max_tokens_per_message=3)
+    message_token_limiter_with_threshold = MessageTokenLimiter(max_tokens_per_message=1, min_tokens=10)
+
+    # Test Parameters
+    message_history_limiter_apply_transform_parameters = {
+        "messages": [long_messages, short_messages, no_content_messages],
+        "expected_messages_len": [3, 3, 2],
+    }
+
+    message_history_limiter_get_logs_parameters = {
+        "messages": [long_messages, short_messages, no_content_messages],
+        "expected_logs": [
+            "Removed 2 messages. Number of messages reduced from 5 to 3.",
+            "No messages were removed.",
+            "No messages were removed.",
+        ],
+        "expected_effect": [True, False, False],
+    }
+
+    message_token_limiter_apply_transform_parameters = {
+        "messages": [long_messages, short_messages, no_content_messages],
+        "expected_token_count": [9, 5, 0],
+        "expected_messages_len": [5, 3, 2],
+    }
+
+    message_token_limiter_with_threshold_apply_transform_parameters = {
+        "messages": [long_messages, short_messages, no_content_messages],
+        "expected_token_count": [5, 5, 0],
+        "expected_messages_len": [5, 3, 2],
+    }
+
+    message_token_limiter_get_logs_parameters = {
+        "messages": [long_messages, short_messages, no_content_messages],
+        "expected_logs": [
+            "Truncated 6 tokens. Number of tokens reduced from 15 to 9",
+            "No tokens were truncated.",
+            "No tokens were truncated.",
+        ],
+        "expected_effect": [True, False, False],
+    }
 
     # Call the MessageHistoryLimiter tests
-    test_message_history_limiter_apply_transform(message_history_limiter, long_messages, 3)
-    test_message_history_limiter_apply_transform(message_history_limiter, short_messages, 3)
-    test_message_history_limiter_get_logs(
-        message_history_limiter, long_messages, "Removed 2 messages. Number of messages reduced from 5 to 3.", True
-    )
-    test_message_history_limiter_get_logs(message_history_limiter, short_messages, "No messages were removed.", False)
+
+    for messages, expected_messages_len in zip(
+        message_history_limiter_apply_transform_parameters["messages"],
+        message_history_limiter_apply_transform_parameters["expected_messages_len"],
+    ):
+        test_message_history_limiter_apply_transform(message_history_limiter, messages, expected_messages_len)
+
+    for messages, expected_logs, expected_effect in zip(
+        message_history_limiter_get_logs_parameters["messages"],
+        message_history_limiter_get_logs_parameters["expected_logs"],
+        message_history_limiter_get_logs_parameters["expected_effect"],
+    ):
+        test_message_history_limiter_get_logs(message_history_limiter, messages, expected_logs, expected_effect)
 
     # Call the MessageTokenLimiter tests
-    test_message_token_limiter_apply_transform(message_token_limiter, long_messages, 9)
-    test_message_token_limiter_apply_transform(message_token_limiter, short_messages, 3)
-    test_message_token_limiter_get_logs(
-        message_token_limiter, long_messages, "Truncated 6 tokens. Number of tokens reduced from 15 to 9", True
-    )
-    test_message_token_limiter_get_logs(message_token_limiter, short_messages, "No tokens were truncated.", False)
+
+    for messages, expected_token_count, expected_messages_len in zip(
+        message_token_limiter_apply_transform_parameters["messages"],
+        message_token_limiter_apply_transform_parameters["expected_token_count"],
+        message_token_limiter_apply_transform_parameters["expected_messages_len"],
+    ):
+        test_message_token_limiter_apply_transform(
+            message_token_limiter, messages, expected_token_count, expected_messages_len
+        )
+
+    for messages, expected_token_count, expected_messages_len in zip(
+        message_token_limiter_with_threshold_apply_transform_parameters["messages"],
+        message_token_limiter_with_threshold_apply_transform_parameters["expected_token_count"],
+        message_token_limiter_with_threshold_apply_transform_parameters["expected_messages_len"],
+    ):
+        test_message_token_limiter_with_threshold_apply_transform(
+            message_token_limiter_with_threshold, messages, expected_token_count, expected_messages_len
+        )
+
+    for messages, expected_logs, expected_effect in zip(
+        message_token_limiter_get_logs_parameters["messages"],
+        message_token_limiter_get_logs_parameters["expected_logs"],
+        message_token_limiter_get_logs_parameters["expected_effect"],
+    ):
+        test_message_token_limiter_get_logs(message_token_limiter, messages, expected_logs, expected_effect)
diff --git a/website/docs/topics/long_contexts.md b/website/docs/topics/long_contexts.md
index 0d8676191044..61e16786d38c 100644
--- a/website/docs/topics/long_contexts.md
+++ b/website/docs/topics/long_contexts.md
@@ -62,11 +62,11 @@ By applying the `MessageHistoryLimiter`, we can see that we were able to limit t
 
 #### Example 2: Limiting the Number of Tokens
 
-To adhere to token limitations, use the `MessageTokenLimiter` transformation. This limits tokens per message and the total token count across all messages:
+To adhere to token limitations, use the `MessageTokenLimiter` transformation. This limits tokens per message and the total token count across all messages. Additionally, a min tokens threshold can be applied:
 
 ```python
 # Limit the token limit per message to 3 tokens
-token_limit_transform = transforms.MessageTokenLimiter(max_tokens_per_message=3)
+token_limit_transform = transforms.MessageTokenLimiter(max_tokens_per_message=3, min_tokens=10)
 
 processed_messages = token_limit_transform.apply_transform(copy.deepcopy(messages))
 
@@ -159,7 +159,7 @@ Now let's add the `TransformMessages` capability to the assistant and run the sa
 context_handling = transform_messages.TransformMessages(
     transforms=[
         transforms.MessageHistoryLimiter(max_messages=10),
-        transforms.MessageTokenLimiter(max_tokens=1000, max_tokens_per_message=50),
+        transforms.MessageTokenLimiter(max_tokens=1000, max_tokens_per_message=50, min_tokens=500),
     ]
 )
 context_handling.add_to_agent(assistant)


Code review comments:
diff_hunk: @@ -194,6 +205,19 @@ def get_logs(self, pre_transform_messages: List[Dict], post_transform_messages:
             return logs_str, True
         return "No tokens were truncated.", False
 
+    def _check_tokens_threshold(self, messages: List[Dict]) -> bool:
path: autogen/agentchat/contrib/capabilities/transforms.py
text: @user1:
we should probably rename this to something more descriptive. If we use it like `if not self._check_tokens_threshold(messages):` it reads out like `if we should not be checking for threshold`. Maybe something like `should_truncate`? `if not self._check_tokens_threshold(messages):` will be read as `if we should start truncating`

@author:
I renamed to `_are_min_tokens_reached`, as a truncation may not be appied even if this function returns `True`.

diff_hunk: @@ -244,6 +268,15 @@ def _validate_max_tokens(self, max_tokens: Optional[int] = None) -> Optional[int
 
         return max_tokens if max_tokens is not None else sys.maxsize
 
+    def _validate_min_tokens(self, min_tokens: int, max_tokens: int) -> int:
+        if min_tokens is None:
+            return 0
+        if min_tokens < 0:
+            raise ValueError("min_tokens must be None or greater than or equal to 0")
+        if max_tokens is not None and min_tokens >= max_tokens:
+            raise ValueError("min_tokens must be less than max_tokens")
path: autogen/agentchat/contrib/capabilities/transforms.py
text: @user1:
```suggestion
            raise ValueError("min_tokens must be less than or equal to max_tokens.")
```

diff_hunk: @@ -244,6 +268,15 @@ def _validate_max_tokens(self, max_tokens: Optional[int] = None) -> Optional[int
 
         return max_tokens if max_tokens is not None else sys.maxsize
 
+    def _validate_min_tokens(self, min_tokens: int, max_tokens: int) -> int:
+        if min_tokens is None:
+            return 0
+        if min_tokens < 0:
+            raise ValueError("min_tokens must be None or greater than or equal to 0")
path: autogen/agentchat/contrib/capabilities/transforms.py
text: @user1:
```suggestion
            raise ValueError("min_tokens must be None or greater than or equal to 0.")
```