diff --git a/lib/crewai/src/crewai/utilities/agent_utils.py b/lib/crewai/src/crewai/utilities/agent_utils.py index b2ab2fc33..3d11ff290 100644 --- a/lib/crewai/src/crewai/utilities/agent_utils.py +++ b/lib/crewai/src/crewai/utilities/agent_utils.py @@ -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) diff --git a/lib/crewai/tests/utilities/test_agent_utils.py b/lib/crewai/tests/utilities/test_agent_utils.py index c4a77006e..ffca78c63 100644 --- a/lib/crewai/tests/utilities/test_agent_utils.py +++ b/lib/crewai/tests/utilities/test_agent_utils.py @@ -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: