From eac81f8693078b84a578f201a70171141558c64f Mon Sep 17 00:00:00 2001 From: Mads Hvelplund Date: Fri, 19 Jun 2026 12:22:52 +0200 Subject: [PATCH] feat(core): MapReduce LLM compression for large exec_command outputs Replaces plain truncation with a parallel summarisation pipeline when exec_command output exceeds STRIX_MAX_TOOL_OUTPUT_CHARS (default 65536). Output is split at JSON record or line boundaries; each chunk is summarised via litellm.acompletion; summaries are consolidated into a single result that fits the context window. Falls back to truncate_exec_result on any compression error. Closes #579 --- strix/agents/factory.py | 45 ++++- strix/core/mapreduce_output.py | 115 +++++++++++++ tests/core/test_mapreduce_output.py | 253 ++++++++++++++++++++++++++++ 3 files changed, 412 insertions(+), 1 deletion(-) create mode 100644 strix/core/mapreduce_output.py create mode 100644 tests/core/test_mapreduce_output.py diff --git a/strix/agents/factory.py b/strix/agents/factory.py index 010bfb40..fca8be03 100644 --- a/strix/agents/factory.py +++ b/strix/agents/factory.py @@ -18,6 +18,7 @@ from pydantic import ValidationError from strix.agents.prompt import render_system_prompt from strix.config import load_settings from strix.core.large_output import truncate_exec_result +from strix.core.mapreduce_output import compress_exec_result from strix.tools.agents_graph.tools import ( agent_finish, create_agent, @@ -207,6 +208,38 @@ def _format_validation_error(tool_name: str, exc: ValidationError) -> str: return f"{tool_name}: invalid arguments — " + "; ".join(parts) +def _resolve_model(ctx: Any) -> str: + """Return the active LLM model string for MapReduce summarisation. + + Prefers the run-context model (set per-agent); falls back to the global + STRIX_LLM setting. Raises ``RuntimeError`` when neither is configured. + """ + run_config = getattr(ctx, "run_config", None) + if run_config is not None: + m = getattr(run_config, "model", None) + if isinstance(m, str) and m.strip(): + return m.strip() + settings_model = load_settings().llm.model + if settings_model and settings_model.strip(): + return settings_model.strip() + raise RuntimeError( + "No LLM model configured for MapReduce compression. " + "Set STRIX_LLM or pass a model to the agent factory." + ) + + +def _extract_task_hint(raw_input: str) -> str: + """Extract the command string from the exec_command JSON payload as a task hint.""" + try: + parsed = json.loads(raw_input) + except (json.JSONDecodeError, TypeError): + return "" + if isinstance(parsed, dict): + cmd = parsed.get("cmd") + return cmd if isinstance(cmd, str) else "" + return "" + + def _wrap_exec_command(tool: FunctionTool) -> FunctionTool: invoke_tool = tool.on_invoke_tool @@ -224,7 +257,17 @@ def _wrap_exec_command(tool: FunctionTool) -> FunctionTool: ) if isinstance(result, str): threshold = load_settings().runtime.max_tool_output_chars - result = truncate_exec_result(result, threshold=threshold) + if threshold > 0 and len(result) > threshold: + try: + model = _resolve_model(ctx) + task_hint = _extract_task_hint(raw_input) + result = await compress_exec_result(result, model=model, task_hint=task_hint) + except Exception: # noqa: BLE001 + logger.warning( + "MapReduce compression failed; falling back to truncation", + exc_info=True, + ) + result = truncate_exec_result(result, threshold=threshold) return result tool.on_invoke_tool = invoke diff --git a/strix/core/mapreduce_output.py b/strix/core/mapreduce_output.py new file mode 100644 index 00000000..2c416998 --- /dev/null +++ b/strix/core/mapreduce_output.py @@ -0,0 +1,115 @@ +"""MapReduce LLM compression for oversized exec_command outputs. + +Splits large output into chunks, summarises each with a lightweight +litellm call in parallel, and consolidates the summaries. Individual +chunk failures produce inline placeholders rather than aborting. +""" + +from __future__ import annotations + +import asyncio +import json +import logging + +import litellm + +from strix.config import load_settings + + +logger = logging.getLogger(__name__) + +CHUNK_SIZE_RECORDS = 100 # JSON records per chunk (~100 KB for typical semgrep findings) +CHUNK_SIZE_LINES = 500 # plain-text lines per chunk +_SUMMARISE_TIMEOUT = 60.0 # per-chunk wall-clock cap (seconds) +_SDK_OUTPUT_SEPARATOR = "\nOutput:\n" + + +async def compress_exec_result(result: str, *, model: str, task_hint: str) -> str: + """Compress an exec_command result produced by the agents SDK. + + Strips the SDK metadata header (Chunk ID / Wall time / exit code), + compresses the content section, then re-attaches the header. + """ + idx = result.find(_SDK_OUTPUT_SEPARATOR) + if idx < 0: + return await compress_large_output(result, model=model, task_hint=task_hint) + sep_end = idx + len(_SDK_OUTPUT_SEPARATOR) + compressed = await compress_large_output(result[sep_end:], model=model, task_hint=task_hint) + return result[:sep_end] + compressed + + +async def compress_large_output(output: str, *, model: str, task_hint: str) -> str: + """Fan *output* into N chunks, summarise each, return consolidated summaries. + + Returns *output* unchanged when it fits in a single chunk (nothing to + fan out). If the consolidated result still exceeds the threshold it is + hard-truncated with an inline warning. + """ + chunks = _split(output) + if len(chunks) <= 1: + return output + summaries = list(await asyncio.gather(*[_summarise(c, model, task_hint) for c in chunks])) + consolidated = _consolidate(summaries, total_chars=len(output), total_chunks=len(chunks)) + threshold = load_settings().runtime.max_tool_output_chars + if threshold > 0 and len(consolidated) > threshold: + return consolidated[:threshold] + "\n[consolidated summary still too large — truncated]" + return consolidated + + +def _split(output: str) -> list[str]: + """Split at JSON record boundaries; fall back to line-count split.""" + stripped = output.lstrip() + if stripped.startswith("["): + try: + records = json.loads(stripped) + if isinstance(records, list): + chunks = [ + json.dumps(records[i : i + CHUNK_SIZE_RECORDS]) + for i in range(0, len(records), CHUNK_SIZE_RECORDS) + ] + return chunks or [output] + except json.JSONDecodeError: + pass + lines = output.splitlines() + if not lines: + return [output] + return [ + "\n".join(lines[i : i + CHUNK_SIZE_LINES]) for i in range(0, len(lines), CHUNK_SIZE_LINES) + ] + + +async def _summarise(chunk: str, model: str, task_hint: str) -> str: + """Summarise one chunk; return a fallback excerpt on any failure.""" + try: + resp = await asyncio.wait_for( + litellm.acompletion( + model=model, + messages=[ + { + "role": "user", + "content": ( + f"Context: {task_hint}\n\n" + "Summarise the key security findings concisely " + "(tool names, severities, affected locations):\n\n" + f"{chunk}" + ), + } + ], + max_tokens=500, + ), + timeout=_SUMMARISE_TIMEOUT, + ) + return str(resp.choices[0].message.content or "[empty]") + except Exception: + logger.exception("Chunk summarisation failed") + return f"[summarisation failed — excerpt:]\n{chunk[:500]}" + + +def _consolidate(summaries: list[str], *, total_chars: int, total_chunks: int) -> str: + """Join summaries under a descriptive header.""" + header = ( + f"[MapReduce compression: {total_chars:,} chars," + f" {total_chunks} chunks → {len(summaries)} summaries]\n\n" + ) + body = "\n\n".join(f"[Chunk {i + 1}/{total_chunks}]\n{s}" for i, s in enumerate(summaries)) + return header + body diff --git a/tests/core/test_mapreduce_output.py b/tests/core/test_mapreduce_output.py new file mode 100644 index 00000000..31f8a8eb --- /dev/null +++ b/tests/core/test_mapreduce_output.py @@ -0,0 +1,253 @@ +"""Tests for strix/core/mapreduce_output.py — MapReduce LLM compression.""" + +from __future__ import annotations + +import json +from typing import TYPE_CHECKING +from unittest.mock import AsyncMock, MagicMock, patch + +from strix.config import load_settings +from strix.core.mapreduce_output import ( + CHUNK_SIZE_LINES, + CHUNK_SIZE_RECORDS, + _consolidate, + _split, + _summarise, + compress_exec_result, + compress_large_output, +) + + +if TYPE_CHECKING: + import pytest + + +# --------------------------------------------------------------------------- +# _split +# --------------------------------------------------------------------------- + + +def test_split_json_array_into_multiple_chunks() -> None: + """250 records at CHUNK_SIZE_RECORDS=100 → 3 chunks, each valid JSON.""" + records = [{"id": i} for i in range(250)] + chunks = _split(json.dumps(records)) + assert len(chunks) == 3 + for chunk in chunks: + assert isinstance(json.loads(chunk), list) + + +def test_split_json_respects_record_boundaries() -> None: + """No record is split across chunk boundaries.""" + records = [{"id": i, "data": "x" * 10} for i in range(CHUNK_SIZE_RECORDS + 1)] + chunks = _split(json.dumps(records)) + assert len(chunks) == 2 + first = json.loads(chunks[0]) + assert len(first) == CHUNK_SIZE_RECORDS + + +def test_split_plain_text_into_chunks() -> None: + """1200 lines at CHUNK_SIZE_LINES=500 → 3 chunks.""" + text = "\n".join(f"line {i}" for i in range(1200)) + chunks = _split(text) + assert len(chunks) == 3 + + +def test_split_invalid_json_falls_back_to_line_split() -> None: + """'[not valid json' triggers line-based split.""" + text = "\n".join("[not json" for _ in range(CHUNK_SIZE_LINES + 1)) + chunks = _split(text) + assert len(chunks) == 2 + + +def test_split_single_chunk_not_split() -> None: + """50 records → 1 chunk (≤ CHUNK_SIZE_RECORDS).""" + records = [{"id": i} for i in range(50)] + chunks = _split(json.dumps(records)) + assert len(chunks) == 1 + + +def test_split_empty_string_returns_one_chunk() -> None: + chunks = _split("") + assert len(chunks) == 1 + + +def test_split_empty_json_array_returns_one_chunk() -> None: + """An empty JSON array [] must return [output], not an empty list.""" + chunks = _split("[]") + assert len(chunks) == 1 + + +def test_split_whitespace_only_returns_one_chunk() -> None: + """Whitespace-only input falls through to line split with one chunk.""" + chunks = _split(" ") + assert len(chunks) == 1 + + +# --------------------------------------------------------------------------- +# _consolidate +# --------------------------------------------------------------------------- + + +def test_consolidate_produces_header() -> None: + result = _consolidate(["a", "b"], total_chars=1000, total_chunks=2) + assert "MapReduce compression" in result + assert "1,000 chars" in result + + +def test_consolidate_numbers_chunks() -> None: + result = _consolidate(["summary0", "summary1"], total_chars=500, total_chunks=2) + assert "[Chunk 1/2]" in result + assert "[Chunk 2/2]" in result + assert "summary0" in result + assert "summary1" in result + + +# --------------------------------------------------------------------------- +# _summarise +# --------------------------------------------------------------------------- + + +async def test_summarise_returns_llm_content() -> None: + mock_resp = MagicMock() + mock_resp.choices[0].message.content = "3 SQL injection findings" + with patch( + "strix.core.mapreduce_output.litellm.acompletion", new_callable=AsyncMock + ) as mock_ac: + mock_ac.return_value = mock_resp + result = await _summarise('["finding"]', "claude-sonnet-4-6", "SQL injection scan") + assert result == "3 SQL injection findings" + + +async def test_summarise_returns_placeholder_on_exception() -> None: + with patch( + "strix.core.mapreduce_output.litellm.acompletion", new_callable=AsyncMock + ) as mock_ac: + mock_ac.side_effect = RuntimeError("API error") + result = await _summarise("chunk_data", "model", "hint") + assert "[summarisation failed" in result + assert "chunk_data"[:100] in result + + +async def test_summarise_returns_placeholder_on_timeout() -> None: + with patch( + "strix.core.mapreduce_output.litellm.acompletion", new_callable=AsyncMock + ) as mock_ac: + mock_ac.side_effect = TimeoutError() + result = await _summarise("chunk_data", "model", "hint") + assert "[summarisation failed" in result + + +# --------------------------------------------------------------------------- +# compress_large_output +# --------------------------------------------------------------------------- + + +async def test_compress_single_chunk_returns_output_unchanged() -> None: + """≤ CHUNK_SIZE_RECORDS records → 1 chunk → returned unchanged, no LLM call.""" + records = [{"id": i} for i in range(50)] + output = json.dumps(records) + with patch( + "strix.core.mapreduce_output.litellm.acompletion", new_callable=AsyncMock + ) as mock_ac: + result = await compress_large_output(output, model="model", task_hint="hint") + mock_ac.assert_not_called() + assert result == output + + +async def test_compress_multi_chunk_consolidates_summaries() -> None: + """250 records → 3 chunks → 3 LLM calls → consolidated output.""" + records = [{"id": i} for i in range(250)] + output = json.dumps(records) + mock_resp = MagicMock() + mock_resp.choices[0].message.content = "summary" + with patch( + "strix.core.mapreduce_output.litellm.acompletion", new_callable=AsyncMock + ) as mock_ac: + mock_ac.return_value = mock_resp + result = await compress_large_output(output, model="model", task_hint="hint") + assert mock_ac.call_count == 3 + assert "MapReduce compression" in result + assert "summary" in result + + +async def test_compress_partial_failure_returns_partial_result() -> None: + """If one chunk fails, its placeholder appears alongside successful summaries.""" + records = [{"id": i} for i in range(250)] + output = json.dumps(records) + good_resp = MagicMock() + good_resp.choices[0].message.content = "good summary" + + call_count = 0 + + async def side_effect(*_args: object, **_kwargs: object) -> MagicMock: + nonlocal call_count + call_count += 1 + if call_count == 2: + raise RuntimeError("chunk 2 failed") + return good_resp + + with patch("strix.core.mapreduce_output.litellm.acompletion", side_effect=side_effect): + result = await compress_large_output(output, model="model", task_hint="hint") + + assert "good summary" in result + assert "[summarisation failed" in result + + +async def test_compress_consolidated_too_large_is_truncated( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """If consolidated result exceeds threshold, it is hard-truncated with a warning.""" + records = [{"id": i} for i in range(250)] + output = json.dumps(records) + mock_resp = MagicMock() + mock_resp.choices[0].message.content = "x" * 300 # 3 summaries x 300 chars = 900+ chars + + settings = load_settings() + monkeypatch.setattr(settings.runtime, "max_tool_output_chars", 100) + + with ( + patch("strix.core.mapreduce_output.load_settings", return_value=settings), + patch("strix.core.mapreduce_output.litellm.acompletion", new_callable=AsyncMock) as mock_ac, + ): + mock_ac.return_value = mock_resp + result = await compress_large_output(output, model="model", task_hint="hint") + + assert len(result) <= 100 + len("\n[consolidated summary still too large — truncated]") + assert "truncated" in result + + +# --------------------------------------------------------------------------- +# compress_exec_result (SDK header handling) +# --------------------------------------------------------------------------- + + +async def test_compress_exec_result_preserves_sdk_header() -> None: + """SDK header up to '\\nOutput:\\n' is returned verbatim.""" + records = [{"id": i} for i in range(250)] + sdk_output = ( + "Chunk ID: abc123\nWall time: 1.23s\nProcess exited with code 0\nOutput:\n" + + json.dumps(records) + ) + mock_resp = MagicMock() + mock_resp.choices[0].message.content = "summary" + with patch( + "strix.core.mapreduce_output.litellm.acompletion", new_callable=AsyncMock + ) as mock_ac: + mock_ac.return_value = mock_resp + result = await compress_exec_result(sdk_output, model="model", task_hint="hint") + assert result.startswith("Chunk ID: abc123") + assert "MapReduce compression" in result + + +async def test_compress_exec_result_no_separator_passes_through() -> None: + """Without SDK separator, entire string is passed to compress_large_output.""" + records = [{"id": i} for i in range(250)] + output = json.dumps(records) + mock_resp = MagicMock() + mock_resp.choices[0].message.content = "summary" + with patch( + "strix.core.mapreduce_output.litellm.acompletion", new_callable=AsyncMock + ) as mock_ac: + mock_ac.return_value = mock_resp + result = await compress_exec_result(output, model="model", task_hint="hint") + assert "MapReduce compression" in result