mirror of
https://github.com/usestrix/strix.git
synced 2026-10-12 03:47:56 +00:00
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
This commit is contained in:
Mads Hvelplund
parent
5aa61c5202
commit
eac81f8693
3 changed files with 412 additions and 1 deletions
|
|
@ -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
|
||||
|
|
|
|||
115
strix/core/mapreduce_output.py
Normal file
115
strix/core/mapreduce_output.py
Normal file
|
|
@ -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
|
||||
253
tests/core/test_mapreduce_output.py
Normal file
253
tests/core/test_mapreduce_output.py
Normal file
|
|
@ -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
|
||||
Loading…
Add table
Reference in a new issue