diff --git a/raven/context_engine/history_trimmer.py b/raven/context_engine/history_trimmer.py index 0c1b2bb..d40d08b 100644 --- a/raven/context_engine/history_trimmer.py +++ b/raven/context_engine/history_trimmer.py @@ -31,8 +31,18 @@ # Provider-safe message keys. Anything else on a session message # (timestamps, internal ids, manifest annotations) is dropped before -# the dict reaches the LLM. -_ALLOWED_KEYS = {"role", "content", "tool_calls", "tool_call_id", "name"} +# the dict reaches the LLM. reasoning_content / thinking_blocks must survive +# so multi-turn reasoning contracts (e.g. DeepSeek thinking mode) hold; the +# provider gate strips thinking_blocks for non-Anthropic targets downstream. +_ALLOWED_KEYS = { + "role", + "content", + "tool_calls", + "tool_call_id", + "name", + "reasoning_content", + "thinking_blocks", +} @dataclass diff --git a/raven/context_engine/segments/curator.py b/raven/context_engine/segments/curator.py index d9e6a34..f0dc723 100644 --- a/raven/context_engine/segments/curator.py +++ b/raven/context_engine/segments/curator.py @@ -293,7 +293,7 @@ def _system_prompt() -> str: @staticmethod def _history_from_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: - allowed = {"role", "content", "tool_calls", "tool_call_id", "name"} + allowed = {"role", "content", "tool_calls", "tool_call_id", "name", "reasoning_content", "thinking_blocks"} out: list[dict[str, Any]] = [] for message in messages: entry = {k: v for k, v in message.items() if k in allowed} diff --git a/raven/utils/helpers.py b/raven/utils/helpers.py index 73355f9..cbf0afc 100644 --- a/raven/utils/helpers.py +++ b/raven/utils/helpers.py @@ -117,6 +117,11 @@ def estimate_prompt_tokens( parts.append(value) if msg.get("tool_calls"): parts.append(json.dumps(msg["tool_calls"], ensure_ascii=False)) + reasoning = msg.get("reasoning_content") + if isinstance(reasoning, str) and reasoning: + parts.append(reasoning) + if msg.get("thinking_blocks"): + parts.append(json.dumps(msg["thinking_blocks"], ensure_ascii=False)) if tools: parts.append(json.dumps(tools, ensure_ascii=False)) @@ -154,6 +159,11 @@ def estimate_message_tokens(message: dict[str, Any]) -> int: parts.append(value) if message.get("tool_calls"): parts.append(json.dumps(message["tool_calls"], ensure_ascii=False)) + reasoning = message.get("reasoning_content") + if isinstance(reasoning, str) and reasoning: + parts.append(reasoning) + if message.get("thinking_blocks"): + parts.append(json.dumps(message["thinking_blocks"], ensure_ascii=False)) payload = "\n".join(parts) if not payload: diff --git a/tests/test_curator_context_engine.py b/tests/test_curator_context_engine.py index 0b3907a..153a1e5 100644 --- a/tests/test_curator_context_engine.py +++ b/tests/test_curator_context_engine.py @@ -8,6 +8,7 @@ from raven.agent.loop import AgentLoop from raven.config import ContextConfig from raven.context_engine import ContextAssembler, TurnContext +from raven.context_engine.segments.curator import CuratorSegmentBuilder from raven.memory_engine.base import TokenBudget from raven.providers.base import LLMProvider, LLMResponse, ToolCallRequest from raven.spine.message import ChatType, Source @@ -195,3 +196,20 @@ async def test_process_message_records_main_and_curator_trajectories(tmp_path: P trace_text = traces[0].read_text(encoding="utf-8") assert "curator_llm_request" in trace_text assert "main_agent_result" in trace_text + + +def test_history_from_messages_preserves_reasoning_fields(): + messages = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": "answer", + "reasoning_content": "chain of thought", + "thinking_blocks": [{"thinking": "block"}], + }, + ] + + history = CuratorSegmentBuilder._history_from_messages(messages) + + assert history[1]["reasoning_content"] == "chain of thought" + assert history[1]["thinking_blocks"] == [{"thinking": "block"}] diff --git a/tests/test_history_trimmer.py b/tests/test_history_trimmer.py new file mode 100644 index 0000000..6d63963 --- /dev/null +++ b/tests/test_history_trimmer.py @@ -0,0 +1,30 @@ +from __future__ import annotations + +from raven.context_engine.history_trimmer import HistoryTrimmer + + +def test_history_from_ids_preserves_reasoning_fields(): + messages = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": "answer", + "reasoning_content": "chain of thought", + "thinking_blocks": [{"thinking": "block"}], + }, + ] + + history = HistoryTrimmer.history_from_ids(messages, [0, 1]) + + assert history[1]["reasoning_content"] == "chain of thought" + assert history[1]["thinking_blocks"] == [{"thinking": "block"}] + + +def test_history_from_ids_drops_non_provider_keys(): + messages = [ + {"role": "user", "content": "hi", "timestamp": "2026-07-08T00:00:00"}, + ] + + history = HistoryTrimmer.history_from_ids(messages, [0]) + + assert history == [{"role": "user", "content": "hi"}] diff --git a/tests/test_token_estimation.py b/tests/test_token_estimation.py index 32e0a6c..11e0de9 100644 --- a/tests/test_token_estimation.py +++ b/tests/test_token_estimation.py @@ -45,6 +45,27 @@ def test_estimate_prompt_tokens_counts_assistant_tool_calls( assert prompt_tokens > 10_000 +def test_estimate_counts_reasoning_fields( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + helpers.tiktoken, + "get_encoding", + lambda _name: _FakeEncoding(), + ) + base = {"role": "assistant", "content": "answer"} + with_reasoning = { + "role": "assistant", + "content": "answer", + "reasoning_content": "r" * 5_000, + "thinking_blocks": [{"thinking": "t" * 5_000}], + } + + assert helpers.estimate_message_tokens(with_reasoning) > helpers.estimate_message_tokens(base) + assert helpers.estimate_prompt_tokens([with_reasoning]) > helpers.estimate_prompt_tokens([base]) + assert helpers.estimate_message_tokens(with_reasoning) > 10_000 + + def test_estimate_prompt_tokens_fallback_counts_tool_calls( monkeypatch: pytest.MonkeyPatch, ) -> None: