Python Deterministic Token Trimming for Message Truncation

[ryanweiler92]

Description

Updating the truncation strategy for messages exceeding the context window to use deterministic token trimming.

Changes Made

Added the _truncate_by_tokens() method to the OpenAIChatCompletionClient class.

The new truncation strategy works as follows:

1. Encode every message once
2. If over the safe window, drop oldest messages until only one needs trimming
3. Clip tokens from the front of that last message
4. Guaranteed ≤ safe window in a single pass

How to Test

from ai_server import ModelEngine, ServerClient

server_connection = ServerClient(
    base="http://localhost:9090/Monolith/api",
    access_key="your_access_key",
    secret_key="your_secret_key",
)

is_connected = server_connection.connected
print(f"Am I connected to the server? {is_connected}")


model = ModelEngine(
    engine_id="4acbe913-df40-4ac0-b28a-daa5ad91b172",
    insight_id=server_connection.cur_insight,
)


def generate_large_token_string(target_tokens=130000):
    base_text = "This is a test sentence to generate many tokens. "

    large_string = base_text * (target_tokens // 6 + 100)

    print(f"Generated string length: {len(large_string)} characters")
    print(f"Estimated tokens: ~{len(large_string) // 4}")

    return large_string


question = "What weighs more, a pound of feathers or a pound of bricks?"
context = "You are a helpful assistant. " + generate_large_token_string(130000)

try:
    response = model.ask(
        question=question,
        context=context,
        param_dict={"stream": False},
    )
    print(response)
except Exception as e:
    print(f"Error: {e}")

Notes

[github-actions]

@CodiumAI-Agent /describe

[QodoAI-Agent]

Title

Python Deterministic Token Trimming for Message Truncation

---

User description

Description

Updating the truncation strategy for messages exceeding the context window to use deterministic token trimming.

Changes Made

Added the _truncate_by_tokens() method to the OpenAIChatCompletionClient class.

The new truncation strategy works as follows:

1. Encode every message once
2. If over the safe window, drop oldest messages until only one needs trimming
3. Clip tokens from the front of that last message
4. Guaranteed ≤ safe window in a single pass

How to Test

from ai_server import ModelEngine, ServerClient

server_connection = ServerClient(
    base="http://localhost:9090/Monolith/api",
    access_key="your_access_key",
    secret_key="your_secret_key",
)

is_connected = server_connection.connected
print(f"Am I connected to the server? {is_connected}")


model = ModelEngine(
    engine_id="4acbe913-df40-4ac0-b28a-daa5ad91b172",
    insight_id=server_connection.cur_insight,
)


def generate_large_token_string(target_tokens=130000):
    base_text = "This is a test sentence to generate many tokens. "

    large_string = base_text * (target_tokens // 6 + 100)

    print(f"Generated string length: {len(large_string)} characters")
    print(f"Estimated tokens: ~{len(large_string) // 4}")

    return large_string


question = "What weighs more, a pound of feathers or a pound of bricks?"
context = "You are a helpful assistant. " + generate_large_token_string(130000)

try:
    response = model.ask(
        question=question,
        context=context,
        param_dict={"stream": False},
    )
    print(response)
except Exception as e:
    print(f"Error: {e}")

Notes

---

PR Type

Enhancement

---

Description

• Implement deterministic ChatML token trimming

• Add safe encode/decode tokenizer methods

• Extend model limits with new Llama model

• Pass context_window and max tokens to tokenizer

---

Changes walkthrough 📝

	Relevant files
Configuration changes	model_limits.py Add Llama-3.1-8B-Instruct model limits                                      py/genai_client/model_limits.py Added meta-llama/Llama-3.1-8B-Instruct limits Kept fallback FALLBACK_CONFIG unchanged +4/-0      openai_api_inference_server.py Pass token limit args to tokenizer                                              py/genai_client/text_generation/openai_clients/openai_api_inference_server.py Forwarded max_completion_tokens to tokenizer Forwarded context_window to tokenizer init +2/-0
Enhancement	openai_chat_completion_client.py Implement deterministic token trimming logic                          py/genai_client/text_generation/openai_clients/openai_chat_completion_client.py Replaced _truncate_messages with _truncate_by_tokens Tokenize once, drop oldest or trim token start Simplified check_token_limits truncation call +54/-98  huggingface_tokenizer.py Add safe encode/decode tokenizer helpers                                  py/genai_client/tokenizers/huggingface_tokenizer.py Added _safe_encode method without special tokens Added _safe_decode fallback for string tokens +17/-0    openai_tokenizer.py Implement safe encode/decode for OpenAI tokenizer                py/genai_client/tokenizers/openai_tokenizer.py Added _safe_encode for text-to-token conversion Added _safe_decode for token-to-text conversion +12/-0

---

Need help?

Type /help how to ... in the comments thread for any questions about PR-Agent usage.

Check out the documentation for more information.

[github-actions]

@CodiumAI-Agent /review

[github-actions]

@CodiumAI-Agent /improve

[QodoAI-Agent]

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

⏱️ Estimated effort to review: 3 🔵🔵🔵⚪⚪

🧪 No relevant tests

🔒 No security concerns identified

⚡ Recommended focus areas for review Truncation Logic Verify that _truncate_by_tokens correctly computes and slices tokens so that only the oldest messages are dropped first and then exactly trims the remaining message to fit within safe_window, while preserving the original message order and roles. def _truncate_by_tokens( self, messages: List[dict], safe_window: int, keep_system: bool = True, ) -> List[dict]: """ Returns a ChatML history whose **total** token count is ≤ safe_window. Oldest non-system messages are dropped first; when only one message needs trimming we cut tokens from its *start*. """ # --- Tokenise *once* ---------------------------------------- toks_per_msg = [] total = 0 for m in messages: toks = self.tokenizer._safe_encode(m["content"]) toks_per_msg.append(toks) total += len(toks) if total <= safe_window: return messages # nothing to do to_cut = total - safe_window # exact excess keep_flags = [True] * len(messages) # --- Build truncation order --------------------------------- # oldest->newest, but leave system message till the end order = list(range(len(messages))) if keep_system and messages and messages[0]["role"] == "system": order = order[1:] + [0] # --- Drop or trim ------------------------------------------- for idx in order: if to_cut == 0: break toks = toks_per_msg[idx] if len(toks) <= to_cut: # drop whole message keep_flags[idx] = False to_cut -= len(toks) else: # keep tail part of this message toks_per_msg[idx] = toks[-(len(toks) - to_cut) :] to_cut = 0 # --- Re-build ChatML ---------------------------------------- new_messages = [] for keep, m, toks in zip(keep_flags, messages, toks_per_msg): if not keep: continue m = m.copy() m["content"] = self.tokenizer._safe_decode(toks) new_messages.append(m) return new_messages Safe encode/decode Ensure _safe_encode and _safe_decode handle both HF and fallback WordCountTokenizer cases consistently, especially regarding special tokens and the format of returned token lists. def _safe_encode(self, text: str) -> List: """ Encode without special tokens; works for real HF tokenizers and for the fallback WordCountTokenizer. """ try: return self.tokenizer.encode(text, add_special_tokens=False) except TypeError: # WordCountTokenizer accepts no kwarg return self.tokenizer.encode(text) def _safe_decode(self, tokens: List) -> str: # WordCountTokenizer returns str tokens; join them if tokens and isinstance(tokens[0], str): return " ".join(tokens) # HF tokenizer returns ints; use its native decode return self.tokenizer.decode(tokens)

[QodoAI-Agent]

PR Code Suggestions ✨

Category	Suggestion	Impact
General	Use constants for pop keys Replace the hard-coded string keys with the defined constants to avoid typos and keep key usage consistent across the codebase. py/genai_client/text_generation/openai_clients/openai_api_inference_server.py [20-21] -max_completion_tokens=init_args.pop("max_completion_tokens", None), -context_window=init_args.pop("context_window", None), +max_completion_tokens=init_args.pop(MAX_COMPLETION_TOKENS, None), +context_window=init_args.pop(CONTEXT_WINDOW, None), Suggestion importance[1-10]: 5 __ Why: Using defined constants (MAX_COMPLETION_TOKENS, CONTEXT_WINDOW) improves consistency and reduces typo risk.	Low
	Break after trimming Add a break immediately after trimming to exit the loop as soon as the required tokens are removed. py/genai_client/text_generation/openai_clients/openai_chat_completion_client.py [262-265] else: # keep tail part of this message - toks_per_msg[idx] = toks[-(len(toks) - to_cut) :] + toks_per_msg[idx] = toks[-(len(toks) - to_cut):] to_cut = 0 + break Suggestion importance[1-10]: 3 __ Why: Adding break exits the loop immediately after trimming and avoids an extra iteration, though the existing check already breaks next cycle.	Low
Possible issue	Preserve system prompt Ensure the system prompt is never removed when keep_system=True by skipping drop/truncate logic for messages with role "system". py/genai_client/text_generation/openai_clients/openai_chat_completion_client.py [254-265] for idx in order: if to_cut == 0: break + # preserve system prompt if requested + if keep_system and messages[idx]["role"] == "system": + continue toks = toks_per_msg[idx] if len(toks) <= to_cut: # drop whole message keep_flags[idx] = False to_cut -= len(toks) else: # keep tail part of this message - toks_per_msg[idx] = toks[-(len(toks) - to_cut) :] + toks_per_msg[idx] = toks[-(len(toks) - to_cut):] to_cut = 0 Suggestion importance[1-10]: 4 __ Why: Skipping truncation for the system message enforces preserving it but may prevent meeting safe_window when all other messages are dropped.	Low

[github-actions]

@CodiumAI-Agent /update_changelog

[QodoAI-Agent]

Changelog updates: 🔄

2025-05-04 [*][https://github.com//pull/765]

Changed

• Implement deterministic token trimming for message truncation via a new _truncate_by_tokens method
• Add _safe_encode/_safe_decode support in tokenizers and update context window/config handling

to commit the new content to the CHANGELOG.md file, please type:
'/update_changelog --pr_update_changelog.push_changelog_changes=true'