mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-09-20 18:13:49 +00:00
[OSS-128] Handle oversized single messages during chunking (#7014)
* fix(OSS-128): split oversized messages before context chunking Normalize single messages that exceed the token budget before boundary chunking so summarization LLM calls do not replay the same context error. Reserve summarization prompt overhead from the chunk budget for large context windows. * refactor(OSS-128): inline message content text extraction as lambda * refactor(OSS-128): restore _message_content_text as a function Revert the lambda assignment to satisfy ruff E731 and keep the helper readable alongside the LLMMessage content shape. * refactor(OSS-128): drop summarization prompt overhead from chunk budget Use the full context window size for message chunking instead of subtracting a fixed prompt overhead. * fix(OSS-128): preserve LLMMessage fields when splitting oversized content Copy the original message attributes into each sub-message and only replace content when expanding oversized entries for chunking. * test(OSS-128): assert rendered summarization requests fit raw context Verify each chunked summarization payload, including system and instruction overhead, stays within the model limit implied by the 85% context window usage ratio.
This commit is contained in:
Vidit Ostwal
@@ -844,6 +844,59 @@ def _estimate_token_count(text: str) -> int:
|
||||
return len(text) // 4
|
||||
|
||||
|
||||
def _message_content_text(msg: LLMMessage) -> str:
|
||||
"""Return the message content as text for token estimation."""
|
||||
content = msg.get("content")
|
||||
if content is None:
|
||||
return ""
|
||||
return str(content)
|
||||
|
||||
|
||||
def _split_text_by_token_limit(text: str, max_tokens: int) -> list[str]:
|
||||
"""Split text into parts each estimated to fit within max_tokens."""
|
||||
if not text:
|
||||
return []
|
||||
if _estimate_token_count(text) <= max_tokens:
|
||||
return [text]
|
||||
|
||||
# Inverse of _estimate_token_count (len // 4): each slice is at most max_tokens.
|
||||
max_chars = max(1, max_tokens * 4)
|
||||
return [text[i : i + max_chars] for i in range(0, len(text), max_chars)]
|
||||
|
||||
|
||||
def _expand_oversized_message(msg: LLMMessage, max_tokens: int) -> list[LLMMessage]:
|
||||
"""Split a message whose content alone exceeds max_tokens into sub-messages."""
|
||||
msg_text = _message_content_text(msg)
|
||||
if _estimate_token_count(msg_text) <= max_tokens:
|
||||
return [msg]
|
||||
|
||||
# Reserve budget for the [Part i/n] prefix added to each sub-message.
|
||||
body_max_tokens = max(1, max_tokens - 5)
|
||||
parts = _split_text_by_token_limit(msg_text, body_max_tokens)
|
||||
total_parts = len(parts)
|
||||
expanded: list[LLMMessage] = []
|
||||
|
||||
for index, part in enumerate(parts, start=1):
|
||||
part_content = (
|
||||
f"[Part {index}/{total_parts}]\n{part}" if total_parts > 1 else part
|
||||
)
|
||||
expanded.append({**msg, "content": part_content})
|
||||
|
||||
return expanded
|
||||
|
||||
|
||||
def _normalize_messages_for_chunking(
|
||||
messages: list[LLMMessage], max_tokens: int
|
||||
) -> list[LLMMessage]:
|
||||
"""Return non-system messages with oversized entries split to fit max_tokens."""
|
||||
normalized: list[LLMMessage] = []
|
||||
for msg in messages:
|
||||
if msg.get("role") == "system":
|
||||
continue
|
||||
normalized.extend(_expand_oversized_message(msg, max_tokens))
|
||||
return normalized
|
||||
|
||||
|
||||
def _format_messages_for_summary(messages: list[LLMMessage]) -> str:
|
||||
"""Format messages with role labels for summarization.
|
||||
|
||||
@@ -904,8 +957,8 @@ def _split_messages_into_chunks(
|
||||
) -> list[list[LLMMessage]]:
|
||||
"""Split messages into chunks at message boundaries.
|
||||
|
||||
Excludes system messages from chunks. Each chunk stays under
|
||||
max_tokens based on estimated token count.
|
||||
Excludes system messages and expands oversized single messages before
|
||||
chunking. Each chunk stays under max_tokens based on estimated token count.
|
||||
|
||||
Args:
|
||||
messages: List of messages to split.
|
||||
@@ -914,24 +967,16 @@ def _split_messages_into_chunks(
|
||||
Returns:
|
||||
List of message chunks.
|
||||
"""
|
||||
non_system = [m for m in messages if m.get("role") != "system"]
|
||||
if not non_system:
|
||||
normalized = _normalize_messages_for_chunking(messages, max_tokens)
|
||||
if not normalized:
|
||||
return []
|
||||
|
||||
chunks: list[list[LLMMessage]] = []
|
||||
current_chunk: list[LLMMessage] = []
|
||||
current_tokens = 0
|
||||
|
||||
for msg in non_system:
|
||||
content = msg.get("content")
|
||||
if content is None:
|
||||
msg_text = ""
|
||||
elif isinstance(content, list):
|
||||
msg_text = str(content)
|
||||
else:
|
||||
msg_text = str(content)
|
||||
|
||||
msg_tokens = _estimate_token_count(msg_text)
|
||||
for msg in normalized:
|
||||
msg_tokens = _estimate_token_count(_message_content_text(msg))
|
||||
|
||||
if current_chunk and (current_tokens + msg_tokens) > max_tokens:
|
||||
chunks.append(current_chunk)
|
||||
|
||||
@@ -17,12 +17,18 @@ from crewai.hooks.tool_hooks import (
|
||||
register_after_tool_call_hook,
|
||||
)
|
||||
from crewai.tools.base_tool import BaseTool
|
||||
from crewai.llm import CONTEXT_WINDOW_USAGE_RATIO
|
||||
from crewai.utilities.agent_utils import (
|
||||
_asummarize_chunks,
|
||||
_estimate_token_count,
|
||||
_expand_oversized_message,
|
||||
_extract_summary_tags,
|
||||
_format_messages_for_summary,
|
||||
_message_content_text,
|
||||
_normalize_messages_for_chunking,
|
||||
_split_messages_into_chunks,
|
||||
_split_text_by_token_limit,
|
||||
format_message_for_llm,
|
||||
convert_tools_to_openai_schema,
|
||||
execute_single_native_tool_call,
|
||||
extract_tool_call_info,
|
||||
@@ -31,6 +37,26 @@ from crewai.utilities.agent_utils import (
|
||||
parse_tool_call_args,
|
||||
summarize_messages,
|
||||
)
|
||||
from crewai.utilities.i18n import I18N_DEFAULT
|
||||
|
||||
|
||||
def _estimate_summarization_request_tokens(chunk: list[dict[str, Any]]) -> int:
|
||||
"""Estimate tokens for the full summarization LLM request for one chunk."""
|
||||
conversation_text = _format_messages_for_summary(chunk)
|
||||
summarization_messages = [
|
||||
format_message_for_llm(
|
||||
I18N_DEFAULT.slice("summarizer_system_message"), role="system"
|
||||
),
|
||||
format_message_for_llm(
|
||||
I18N_DEFAULT.slice("summarize_instruction").format(
|
||||
conversation=conversation_text
|
||||
),
|
||||
),
|
||||
]
|
||||
return sum(
|
||||
_estimate_token_count(str(message.get("content", "")))
|
||||
for message in summarization_messages
|
||||
)
|
||||
|
||||
|
||||
class CalculatorInput(BaseModel):
|
||||
@@ -673,6 +699,170 @@ class TestSplitMessagesIntoChunks:
|
||||
assert len(chunks) == 1
|
||||
assert len(chunks[0]) == 2
|
||||
|
||||
def test_splits_oversized_single_message(self) -> None:
|
||||
messages: list[dict[str, Any]] = [
|
||||
{"role": "tool", "content": "X" * 1200, "name": "web_scraper"},
|
||||
]
|
||||
max_tokens = 100
|
||||
chunks = _split_messages_into_chunks(messages, max_tokens=max_tokens)
|
||||
assert len(chunks) > 1
|
||||
for chunk in chunks:
|
||||
chunk_tokens = sum(
|
||||
_estimate_token_count(_message_content_text(msg)) for msg in chunk
|
||||
)
|
||||
assert chunk_tokens <= max_tokens
|
||||
|
||||
def test_oversized_tool_in_conversation(self) -> None:
|
||||
messages: list[dict[str, Any]] = [
|
||||
{"role": "user", "content": "Search"},
|
||||
{"role": "tool", "content": "Y" * 1200, "name": "search"},
|
||||
{"role": "assistant", "content": "Done"},
|
||||
]
|
||||
max_tokens = 100
|
||||
chunks = _split_messages_into_chunks(messages, max_tokens=max_tokens)
|
||||
assert len(chunks) > 1
|
||||
for chunk in chunks:
|
||||
chunk_tokens = sum(
|
||||
_estimate_token_count(_message_content_text(msg)) for msg in chunk
|
||||
)
|
||||
assert chunk_tokens <= max_tokens
|
||||
|
||||
def test_rendered_summarization_request_within_raw_context_window(self) -> None:
|
||||
"""Chunked payloads plus summarization prompt fit in the raw model limit.
|
||||
|
||||
get_context_window_size() already applies CONTEXT_WINDOW_USAGE_RATIO (85%).
|
||||
The remaining 15% headroom should absorb summarizer system/instruction overhead.
|
||||
"""
|
||||
chunk_budget = 50_000
|
||||
raw_context_limit = int(chunk_budget / CONTEXT_WINDOW_USAGE_RATIO)
|
||||
messages: list[dict[str, Any]] = [
|
||||
{"role": "user", "content": "Fetch CRM data for the attendee."},
|
||||
{"role": "tool", "content": "Z" * 200_000, "name": "hubspot_search"},
|
||||
{"role": "assistant", "content": "Collected HubSpot results."},
|
||||
]
|
||||
|
||||
chunks = _split_messages_into_chunks(messages, max_tokens=chunk_budget)
|
||||
assert len(chunks) > 1
|
||||
|
||||
for chunk in chunks:
|
||||
request_tokens = _estimate_summarization_request_tokens(chunk)
|
||||
assert request_tokens <= raw_context_limit
|
||||
|
||||
|
||||
class TestMessageContentText:
|
||||
"""Tests for _message_content_text helper."""
|
||||
|
||||
def test_string_content(self) -> None:
|
||||
msg: dict[str, Any] = {"role": "user", "content": "hello"}
|
||||
assert _message_content_text(msg) == "hello"
|
||||
|
||||
def test_none_content(self) -> None:
|
||||
msg: dict[str, Any] = {"role": "assistant", "content": None}
|
||||
assert _message_content_text(msg) == ""
|
||||
|
||||
def test_list_content_uses_str(self) -> None:
|
||||
msg: dict[str, Any] = {
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": "first"}],
|
||||
}
|
||||
assert _message_content_text(msg) == str(msg["content"])
|
||||
|
||||
|
||||
class TestSplitTextByTokenLimit:
|
||||
"""Tests for _split_text_by_token_limit helper."""
|
||||
|
||||
def test_empty_string(self) -> None:
|
||||
assert _split_text_by_token_limit("", max_tokens=100) == []
|
||||
|
||||
def test_under_limit_returns_single_part(self) -> None:
|
||||
assert _split_text_by_token_limit("hello", max_tokens=100) == ["hello"]
|
||||
|
||||
def test_split_preserves_content(self) -> None:
|
||||
text = "a" * 600
|
||||
parts = _split_text_by_token_limit(text, max_tokens=100)
|
||||
assert len(parts) > 1
|
||||
assert "".join(parts) == text
|
||||
|
||||
def test_each_part_estimated_under_limit(self) -> None:
|
||||
text = "b" * 1200
|
||||
max_tokens = 100
|
||||
parts = _split_text_by_token_limit(text, max_tokens=max_tokens)
|
||||
assert all(_estimate_token_count(part) <= max_tokens for part in parts)
|
||||
|
||||
|
||||
class TestExpandOversizedMessage:
|
||||
"""Tests for _expand_oversized_message helper."""
|
||||
|
||||
def test_returns_original_when_under_limit(self) -> None:
|
||||
msg: dict[str, Any] = {"role": "user", "content": "hello"}
|
||||
expanded = _expand_oversized_message(msg, max_tokens=100)
|
||||
assert expanded == [msg]
|
||||
|
||||
def test_splits_tool_output_with_metadata(self) -> None:
|
||||
msg: dict[str, Any] = {
|
||||
"role": "tool",
|
||||
"content": "Z" * 1200,
|
||||
"name": "fetch_page",
|
||||
"tool_call_id": "call_123",
|
||||
}
|
||||
expanded = _expand_oversized_message(msg, max_tokens=100)
|
||||
assert len(expanded) > 1
|
||||
assert all(part["role"] == "tool" for part in expanded)
|
||||
assert all(part["name"] == "fetch_page" for part in expanded)
|
||||
assert all(part["tool_call_id"] == "call_123" for part in expanded)
|
||||
assert expanded[0]["content"].startswith("[Part 1/")
|
||||
|
||||
def test_preserves_non_content_fields(self) -> None:
|
||||
mock_file = MagicMock()
|
||||
msg: dict[str, Any] = {
|
||||
"role": "user",
|
||||
"content": "X" * 1200,
|
||||
"files": {"report.pdf": mock_file},
|
||||
}
|
||||
expanded = _expand_oversized_message(msg, max_tokens=100)
|
||||
assert len(expanded) > 1
|
||||
assert all(part["role"] == "user" for part in expanded)
|
||||
assert all(part["files"] == {"report.pdf": mock_file} for part in expanded)
|
||||
|
||||
def test_each_part_estimated_under_limit(self) -> None:
|
||||
msg: dict[str, Any] = {"role": "user", "content": "Y" * 1200}
|
||||
max_tokens = 100
|
||||
expanded = _expand_oversized_message(msg, max_tokens=max_tokens)
|
||||
assert len(expanded) > 1
|
||||
assert all(
|
||||
_estimate_token_count(_message_content_text(part)) <= max_tokens
|
||||
for part in expanded
|
||||
)
|
||||
|
||||
|
||||
class TestNormalizeMessagesForChunking:
|
||||
"""Tests for _normalize_messages_for_chunking helper."""
|
||||
|
||||
def test_excludes_system_messages(self) -> None:
|
||||
messages: list[dict[str, Any]] = [
|
||||
{"role": "system", "content": "System prompt"},
|
||||
{"role": "user", "content": "Hello"},
|
||||
]
|
||||
normalized = _normalize_messages_for_chunking(messages, max_tokens=1000)
|
||||
assert len(normalized) == 1
|
||||
assert normalized[0]["role"] == "user"
|
||||
|
||||
def test_expands_oversized_and_preserves_small_messages(self) -> None:
|
||||
messages: list[dict[str, Any]] = [
|
||||
{"role": "user", "content": "Short"},
|
||||
{"role": "tool", "content": "X" * 1200, "name": "search"},
|
||||
{"role": "assistant", "content": "Done"},
|
||||
]
|
||||
max_tokens = 100
|
||||
normalized = _normalize_messages_for_chunking(messages, max_tokens=max_tokens)
|
||||
assert normalized[0]["content"] == "Short"
|
||||
assert normalized[-1]["content"] == "Done"
|
||||
assert len(normalized) > 3
|
||||
assert all(
|
||||
_estimate_token_count(_message_content_text(msg)) <= max_tokens
|
||||
for msg in normalized
|
||||
)
|
||||
|
||||
|
||||
class TestEstimateTokenCount:
|
||||
"""Tests for _estimate_token_count helper."""
|
||||
@@ -705,7 +895,9 @@ class TestParallelSummarization:
|
||||
"""
|
||||
msgs: list[dict[str, Any]] = []
|
||||
for i in range(n):
|
||||
msgs.append({"role": "user", "content": f"msg-{i} " + "x" * 400})
|
||||
prefix = f"msg-{i} "
|
||||
padding = "x" * max(0, 400 - len(prefix))
|
||||
msgs.append({"role": "user", "content": prefix + padding})
|
||||
return msgs
|
||||
|
||||
def test_multiple_chunks_use_acall(self) -> None:
|
||||
|
||||
Reference in New Issue
Block a user