diff --git a/raven/context_engine/history_trimmer.py b/raven/context_engine/history_trimmer.py index b4a62ca..2f68d1e 100644 --- a/raven/context_engine/history_trimmer.py +++ b/raven/context_engine/history_trimmer.py @@ -167,6 +167,46 @@ def _first_droppable(ids: list[int], protected_ids: set[int]) -> int | None: return pos return 0 if ids else None + @staticmethod + def _tool_exchange_ids(messages: list[dict[str, Any]], mid: int) -> set[int]: + """Return every message that must be dropped with ``mid``. + + One assistant message may request several tools. Its call message and + every corresponding result therefore form one structural unit: keeping + only part of that unit produces an invalid provider request. + """ + message = messages[mid] + parent_idx: int | None = None + if message.get("role") == "assistant" and message.get("tool_calls"): + parent_idx = mid + elif message.get("role") == "tool" and message.get("tool_call_id"): + call_id = str(message["tool_call_id"]) + for idx, candidate in enumerate(messages): + if candidate.get("role") != "assistant": + continue + if any( + isinstance(tool_call, dict) and str(tool_call.get("id", "")) == call_id + for tool_call in candidate.get("tool_calls") or [] + ): + parent_idx = idx + break + + if parent_idx is None: + return {mid} + + call_ids = { + str(tool_call["id"]) + for tool_call in messages[parent_idx].get("tool_calls") or [] + if isinstance(tool_call, dict) and tool_call.get("id") + } + exchange = {parent_idx} + exchange.update( + idx + for idx, candidate in enumerate(messages) + if candidate.get("role") == "tool" and str(candidate.get("tool_call_id", "")) in call_ids + ) + return exchange + # ------------------------------------------------------------------ # Budget-driven trimming # ------------------------------------------------------------------ @@ -205,8 +245,10 @@ def trim( drop_idx = self._first_droppable(trimmed_ids, protected_ids) if drop_idx is None: break - dropped = trimmed_ids.pop(drop_idx) - warnings.append(f"dropped message {dropped} to fit budget") + dropped = trimmed_ids[drop_idx] + dropped_group = self._tool_exchange_ids(session_messages, dropped) + trimmed_ids = [mid for mid in trimmed_ids if mid not in dropped_group] + warnings.extend(f"dropped message {mid} to fit budget" for mid in sorted(dropped_group)) history = self.history_from_ids(session_messages, trimmed_ids) messages = build_messages(history) estimated, source = estimate_prompt_tokens_chain( diff --git a/tests/test_history_trimmer.py b/tests/test_history_trimmer.py index 6d63963..7166f30 100644 --- a/tests/test_history_trimmer.py +++ b/tests/test_history_trimmer.py @@ -3,6 +3,15 @@ from raven.context_engine.history_trimmer import HistoryTrimmer +class _WeightedProvider: + def estimate_prompt_tokens(self, messages, tools, model): + del tools, model + total = 0 + for message in messages: + total += 100 if message.get("tool_calls") else 1 + return total, "test" + + def test_history_from_ids_preserves_reasoning_fields(): messages = [ {"role": "user", "content": "hi"}, @@ -28,3 +37,42 @@ def test_history_from_ids_drops_non_provider_keys(): history = HistoryTrimmer.history_from_ids(messages, [0]) assert history == [{"role": "user", "content": "hi"}] + + +def test_budget_trim_drops_tool_call_and_result_as_one_group(): + messages = [ + {"role": "user", "content": "keep me"}, + { + "role": "assistant", + "content": "calling a tool", + "tool_calls": [ + { + "id": "call_1|fc_1", + "type": "function", + "function": {"name": "message", "arguments": "{}"}, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_1|fc_1", + "name": "message", + "content": "sent", + }, + ] + trimmer = HistoryTrimmer(_WeightedProvider(), "test-model", lambda: [], 10) + + built, outcome = trimmer.trim( + session_messages=messages, + ids=[0, 1, 2], + protected_ids={0}, + reserved_output=0, + build_messages=lambda history: [ + {"role": "system", "content": "system"}, + *history, + {"role": "user", "content": "current"}, + ], + ) + + assert outcome.included_ids == [0] + assert HistoryTrimmer.structural_errors(built) == []