refactor: remove _nr_to_assistant_message shim + fix flush_memories guard

NormalizedResponse and ToolCall now have backward-compat properties
so the agent loop can read them directly without the shim:

  ToolCall: .type, .function (returns self), .call_id, .response_item_id
  NormalizedResponse: .reasoning_content, .reasoning_details,
                      .codex_reasoning_items

This eliminates the 35-line shim and its 4 call sites in run_agent.py.

Also changes flush_memories guard from hasattr(response, 'choices')
to self.api_mode in ('chat_completions', 'bedrock_converse') so it
works with raw boto3 dicts too.

WS1 items 3+4 of Cycle 2 (#14418).
This commit is contained in:
kshitijk4poor
2026-04-23 14:06:36 +05:30
committed by Teknium
parent f4612785a4
commit 43de1ca8c2
8 changed files with 233 additions and 157 deletions

View File

@@ -18,12 +18,12 @@ from agent.anthropic_adapter import (
convert_messages_to_anthropic,
convert_tools_to_anthropic,
is_claude_code_token_valid,
normalize_anthropic_response,
normalize_model_name,
read_claude_code_credentials,
resolve_anthropic_token,
run_oauth_setup_token,
)
from agent.transports import get_transport
# ---------------------------------------------------------------------------
@@ -1242,7 +1242,7 @@ class TestNormalizeResponse:
def test_text_response(self):
block = SimpleNamespace(type="text", text="Hello world")
nr = normalize_anthropic_response(self._make_response([block]))
nr = get_transport("anthropic_messages").normalize_response(self._make_response([block]))
assert nr.content == "Hello world"
assert nr.finish_reason == "stop"
assert nr.tool_calls is None
@@ -1257,7 +1257,7 @@ class TestNormalizeResponse:
input={"query": "test"},
),
]
nr = normalize_anthropic_response(
nr = get_transport("anthropic_messages").normalize_response(
self._make_response(blocks, "tool_use")
)
assert nr.content == "Searching..."
@@ -1271,7 +1271,7 @@ class TestNormalizeResponse:
SimpleNamespace(type="thinking", thinking="Let me reason about this..."),
SimpleNamespace(type="text", text="The answer is 42."),
]
nr = normalize_anthropic_response(self._make_response(blocks))
nr = get_transport("anthropic_messages").normalize_response(self._make_response(blocks))
assert nr.content == "The answer is 42."
assert nr.reasoning == "Let me reason about this..."
assert nr.provider_data["reasoning_details"] == [{"type": "thinking", "thinking": "Let me reason about this..."}]
@@ -1285,19 +1285,19 @@ class TestNormalizeResponse:
redacted=False,
),
]
nr = normalize_anthropic_response(self._make_response(blocks))
nr = get_transport("anthropic_messages").normalize_response(self._make_response(blocks))
assert nr.provider_data["reasoning_details"][0]["signature"] == "opaque_signature"
assert nr.provider_data["reasoning_details"][0]["thinking"] == "Let me reason about this..."
def test_stop_reason_mapping(self):
block = SimpleNamespace(type="text", text="x")
nr1 = normalize_anthropic_response(
nr1 = get_transport("anthropic_messages").normalize_response(
self._make_response([block], "end_turn")
)
nr2 = normalize_anthropic_response(
nr2 = get_transport("anthropic_messages").normalize_response(
self._make_response([block], "tool_use")
)
nr3 = normalize_anthropic_response(
nr3 = get_transport("anthropic_messages").normalize_response(
self._make_response([block], "max_tokens")
)
assert nr1.finish_reason == "stop"
@@ -1310,10 +1310,10 @@ class TestNormalizeResponse:
# handlers already understand, instead of silently collapsing to
# "stop" (old behavior).
block = SimpleNamespace(type="text", text="")
nr_refusal = normalize_anthropic_response(
nr_refusal = get_transport("anthropic_messages").normalize_response(
self._make_response([block], "refusal")
)
nr_overflow = normalize_anthropic_response(
nr_overflow = get_transport("anthropic_messages").normalize_response(
self._make_response([block], "model_context_window_exceeded")
)
assert nr_refusal.finish_reason == "content_filter"
@@ -1323,7 +1323,7 @@ class TestNormalizeResponse:
block = SimpleNamespace(
type="tool_use", id="tc_1", name="search", input={"q": "hi"}
)
nr = normalize_anthropic_response(
nr = get_transport("anthropic_messages").normalize_response(
self._make_response([block], "tool_use")
)
assert nr.content is None

View File

@@ -149,3 +149,95 @@ class TestMapFinishReason:
def test_none_reason(self):
assert map_finish_reason(None, self.ANTHROPIC_MAP) == "stop"
# ---------------------------------------------------------------------------
# Backward-compat property tests
# ---------------------------------------------------------------------------
class TestToolCallBackwardCompat:
"""Test duck-typing properties that let ToolCall pass through code expecting
the old SimpleNamespace(id, type, function=SimpleNamespace(name, arguments)) shape."""
def test_type_is_function(self):
tc = ToolCall(id="1", name="search", arguments='{"q":"test"}')
assert tc.type == "function"
def test_function_returns_self(self):
tc = ToolCall(id="1", name="search", arguments='{"q":"test"}')
assert tc.function is tc
def test_function_name_matches(self):
tc = ToolCall(id="1", name="search", arguments='{"q":"test"}')
assert tc.function.name == "search"
assert tc.function.name == tc.name
def test_function_arguments_matches(self):
tc = ToolCall(id="1", name="search", arguments='{"q":"test"}')
assert tc.function.arguments == '{"q":"test"}'
assert tc.function.arguments == tc.arguments
def test_call_id_from_provider_data(self):
tc = ToolCall(id="1", name="fn", arguments="{}", provider_data={"call_id": "c1"})
assert tc.call_id == "c1"
def test_call_id_none_when_no_provider_data(self):
tc = ToolCall(id="1", name="fn", arguments="{}", provider_data=None)
assert tc.call_id is None
def test_response_item_id_from_provider_data(self):
tc = ToolCall(id="1", name="fn", arguments="{}", provider_data={"response_item_id": "r1"})
assert tc.response_item_id == "r1"
def test_response_item_id_none_when_missing(self):
tc = ToolCall(id="1", name="fn", arguments="{}", provider_data={"call_id": "c1"})
assert tc.response_item_id is None
def test_getattr_pattern_matches_agent_loop(self):
"""run_agent.py uses getattr(tool_call, 'call_id', None) — verify it works."""
tc = ToolCall(id="1", name="fn", arguments="{}", provider_data={"call_id": "c1"})
assert getattr(tc, "call_id", None) == "c1"
tc_no_pd = ToolCall(id="1", name="fn", arguments="{}")
assert getattr(tc_no_pd, "call_id", None) is None
class TestNormalizedResponseBackwardCompat:
"""Test properties that replaced _nr_to_assistant_message() shim."""
def test_reasoning_content_from_provider_data(self):
nr = NormalizedResponse(
content="hi", tool_calls=None, finish_reason="stop",
provider_data={"reasoning_content": "thought process"},
)
assert nr.reasoning_content == "thought process"
def test_reasoning_content_none_when_absent(self):
nr = NormalizedResponse(content="hi", tool_calls=None, finish_reason="stop")
assert nr.reasoning_content is None
def test_reasoning_details_from_provider_data(self):
details = [{"type": "thinking", "thinking": "hmm"}]
nr = NormalizedResponse(
content="hi", tool_calls=None, finish_reason="stop",
provider_data={"reasoning_details": details},
)
assert nr.reasoning_details == details
def test_reasoning_details_none_when_no_provider_data(self):
nr = NormalizedResponse(
content="hi", tool_calls=None, finish_reason="stop",
provider_data=None,
)
assert nr.reasoning_details is None
def test_codex_reasoning_items_from_provider_data(self):
items = ["item1", "item2"]
nr = NormalizedResponse(
content="hi", tool_calls=None, finish_reason="stop",
provider_data={"codex_reasoning_items": items},
)
assert nr.codex_reasoning_items == items
def test_codex_reasoning_items_none_when_absent(self):
nr = NormalizedResponse(content="hi", tool_calls=None, finish_reason="stop")
assert nr.codex_reasoning_items is None

View File

@@ -47,16 +47,16 @@ def _make_anthropic_response(blocks, stop_reason: str = "max_tokens"):
class TestTruncatedAnthropicResponseNormalization:
"""normalize_anthropic_response() gives us the shape _build_assistant_message expects."""
"""AnthropicTransport.normalize_response() gives us the shape _build_assistant_message expects."""
def test_text_only_truncation_produces_text_content_no_tool_calls(self):
"""Pure-text Anthropic truncation → continuation path should fire."""
from agent.anthropic_adapter import normalize_anthropic_response
from agent.transports import get_transport
response = _make_anthropic_response(
[_make_anthropic_text_block("partial response that was cut off")]
)
nr = normalize_anthropic_response(response)
nr = get_transport("anthropic_messages").normalize_response(response)
# The continuation block checks these two attributes:
# assistant_message.content → appended to truncated_response_prefix
@@ -71,7 +71,7 @@ class TestTruncatedAnthropicResponseNormalization:
def test_truncated_tool_call_produces_tool_calls(self):
"""Tool-use truncation → tool-call retry path should fire."""
from agent.anthropic_adapter import normalize_anthropic_response
from agent.transports import get_transport
response = _make_anthropic_response(
[
@@ -79,7 +79,7 @@ class TestTruncatedAnthropicResponseNormalization:
_make_anthropic_tool_use_block(),
]
)
nr = normalize_anthropic_response(response)
nr = get_transport("anthropic_messages").normalize_response(response)
assert bool(nr.tool_calls), (
"Truncation mid-tool_use must expose tool_calls so the "
@@ -89,10 +89,10 @@ class TestTruncatedAnthropicResponseNormalization:
def test_empty_content_does_not_crash(self):
"""Empty response.content — defensive: treat as a truncation with no text."""
from agent.anthropic_adapter import normalize_anthropic_response
from agent.transports import get_transport
response = _make_anthropic_response([])
nr = normalize_anthropic_response(response)
nr = get_transport("anthropic_messages").normalize_response(response)
# Depending on the adapter, content may be "" or None — both are
# acceptable; what matters is no exception.
assert nr is not None