From 478ef45584a571d940e90d169f7aae8caa234062 Mon Sep 17 00:00:00 2001 From: Andiii208 <246843242+Andiii208@users.noreply.github.com> Date: Sun, 4 Oct 2026 00:07:43 +0800 Subject: [PATCH 1/5] fix(agents): align mode-driven token accounting with prepared requests Addresses review feedback on the mode-consistency fix for issue #282: - Normalize multimodal history back to chat-format content parts for the token counter regardless of the agent's mode. The Responses-format parts (input_image) that item.to_openai(RESPONSES_TOOLS) produces raise TokenCountError in LiteLLM's chat counter, so deriving the mode from a Responses client made get_context_token_count() fail on image history. - Replace the name-suffix mode family with an explicit set of modes whose prepared request carries the schema as a tool/function definition. The previous whitelist (TOOLS, TOOLS_STRICT, PARALLEL_TOOLS) left native provider modes such as ANTHROPIC_TOOLS counting the schema as JSON-schema prose while Instructor's prepared request sends a tool, so accounting silently disagreed with the request. COHERE_TOOLS stays out of the set: despite its name it embeds the schema in an instruction message. - Cover both with tests that count real tokens and that prepare a request per Instructor mode and assert the accounting set matches what is sent. Simplifications: _mode_family is gone (the warning now compares membership in the same set that drives _build_tools_definition()), and the spy-based serialization test is replaced by real counting tests. Co-Authored-By: Claude Code --- .../atomic_agents/agents/atomic_agent.py | 95 ++++++- .../agents/test_agent_mode_consistency.py | 244 ++++++++++++++++++ .../skills/framework/references/providers.md | 9 +- docs/api/agents.md | 1 + docs/guides/quickstart.md | 7 + 5 files changed, 349 insertions(+), 7 deletions(-) create mode 100644 atomic-agents/tests/agents/test_agent_mode_consistency.py diff --git a/atomic-agents/atomic_agents/agents/atomic_agent.py b/atomic-agents/atomic_agents/agents/atomic_agent.py index ab656e7b..72f8ac6c 100644 --- a/atomic-agents/atomic_agents/agents/atomic_agent.py +++ b/atomic-agents/atomic_agents/agents/atomic_agent.py @@ -18,6 +18,39 @@ from instructor.dsl.partial import PartialBase from jiter import from_json +# Instructor modes whose prepared request carries the output schema as a +# tool/function definition instead of text inside the messages. Token accounting +# must agree with the actual request, so these modes are counted with a tools +# definition while all other modes append the schema to the system message. +# Keep aligned with Instructor's per-mode request preparation; the parametrized +# test in tests/agents/test_agent_mode_consistency.py prepares a request per +# mode and asserts this set matches what is actually sent. +# Mode.COHERE_TOOLS is deliberately absent: despite its name, its prepared +# request embeds the schema in an instruction message and sends no tool. +_TOOL_MODES = frozenset( + { + Mode.FUNCTIONS, + Mode.PARALLEL_TOOLS, + Mode.TOOLS, + Mode.TOOLS_STRICT, + Mode.RESPONSES_TOOLS, + Mode.RESPONSES_TOOLS_WITH_INBUILT_TOOLS, + Mode.ANTHROPIC_TOOLS, + Mode.ANTHROPIC_REASONING_TOOLS, + Mode.ANTHROPIC_PARALLEL_TOOLS, + Mode.MISTRAL_TOOLS, + Mode.VERTEXAI_TOOLS, + Mode.VERTEXAI_PARALLEL_TOOLS, + Mode.GEMINI_TOOLS, + Mode.GENAI_TOOLS, + Mode.XAI_TOOLS, + Mode.CEREBRAS_TOOLS, + Mode.FIREWORKS_TOOLS, + Mode.WRITER_TOOLS, + Mode.BEDROCK_TOOLS, + } +) + def model_from_chunks_patched(cls, json_chunks, **kwargs): potential_object = "" @@ -91,7 +124,16 @@ class AgentConfig(BaseModel): ), ) model_config = {"arbitrary_types_allowed": True} - mode: Mode = Field(default=Mode.TOOLS, description="The Instructor mode used for structured outputs (TOOLS, JSON, etc.).") + mode: Optional[Mode] = Field( + default=None, + description=( + "The Instructor mode used for structured outputs (TOOLS, JSON, etc.). " + "Drives token accounting only: modes that transmit the schema as a tool " + "definition count it as a tool, others append it to the system message. " + "When None, the client's mode is used; an explicit mode must agree with " + "the client's mode on how the schema is transmitted, or a warning is logged." + ), + ) model_api_parameters: Optional[dict] = Field(None, description="Additional parameters passed to the API provider.") max_context_tokens: Optional[int] = Field( None, @@ -202,7 +244,7 @@ def __init__(self, config: AgentConfig): self.tool_result_role = "user" if config.assistant_role == "model" else "system" self.initial_history = self.history.copy() self.current_user_input = None - self.mode = config.mode + self.mode = self._resolve_mode(config) self.model_api_parameters = config.model_api_parameters or {} self.max_context_tokens = config.max_context_tokens @@ -216,6 +258,45 @@ def reset_history(self): """ self.history = self.initial_history.copy() + @staticmethod + def _resolve_mode(config: AgentConfig) -> Mode: + """ + Resolves the effective Instructor mode used for token accounting. + + The API call format is decided by the mode the Instructor client was created + with, so `AgentConfig.mode` follows the client by default. An explicitly + configured mode that disagrees with the client on how the schema is + transmitted warns, because token accounting then silently disagrees with + the actual API calls. + + Args: + config (AgentConfig): Configuration for the chat agent. + + Returns: + Mode: The mode to use for token accounting. + """ + client_mode = getattr(config.client, "mode", None) + if not isinstance(client_mode, Mode): + # Clients that do not expose a Mode (or test doubles of them) cannot be + # compared; the explicitly configured mode, if any, is used as-is. + client_mode = None + + if config.mode is None: + return client_mode if client_mode is not None else Mode.TOOLS + + if client_mode is not None and (config.mode in _TOOL_MODES) != (client_mode in _TOOL_MODES): + logger = logging.getLogger(__name__) + logger.warning( + "AgentConfig.mode (%s) and the Instructor client's mode (%s) disagree on how the " + "output schema is transmitted, so token accounting will silently disagree with " + "the actual API calls. Create the Instructor client with %s as well, or omit " + "AgentConfig.mode to follow the client automatically.", + config.mode.name, + client_mode.name, + config.mode.name, + ) + return config.mode + def add_tool_result(self, content: BaseIOSchema) -> None: """ Adds a tool result or context injection to the chat history using the @@ -391,9 +472,8 @@ def _build_tools_definition(self) -> Optional[List[Dict[str, Any]]]: """ from instructor.processing.schema import generate_openai_schema - # Only return tools for TOOLS-based modes - tools_modes = {Mode.TOOLS, Mode.TOOLS_STRICT, Mode.PARALLEL_TOOLS} - if self.mode in tools_modes: + # Only return tools for modes that transmit the schema as a tool definition + if self.mode in _TOOL_MODES: return [ { "type": "function", @@ -451,7 +531,10 @@ def _serialize_history_for_token_count(self) -> List[Dict[str, Any]]: # Text content - wrap in OpenAI text format serialized_content.append({"type": "text", "text": item}) elif isinstance(item, (Image, Audio, PDF)): - # Multimodal object - use instructor's to_openai method + # Multimodal object - serialize with Instructor. The token + # counter only accepts chat-format content parts, so media is + # normalized to that format regardless of the agent's mode + # (Responses-format parts such as input_image raise). try: serialized_content.append(item.to_openai(Mode.JSON)) except Exception as e: diff --git a/atomic-agents/tests/agents/test_agent_mode_consistency.py b/atomic-agents/tests/agents/test_agent_mode_consistency.py new file mode 100644 index 00000000..99168a96 --- /dev/null +++ b/atomic-agents/tests/agents/test_agent_mode_consistency.py @@ -0,0 +1,244 @@ +"""Tests for AgentConfig mode resolution and client/config mode consistency (issue #282). + +The mode on the Instructor client decides the API call format, while +``AgentConfig.mode`` only drives token accounting. These tests cover deriving the +mode from the client, warning when an explicit mode disagrees with the client on +how the schema is transmitted, and keeping the accounting aligned with the +requests Instructor actually prepares. +""" + +import logging +from unittest.mock import Mock + +import instructor +import pytest +from instructor.processing.response import handle_response_model +from openai import OpenAI +from pydantic import Field + +from atomic_agents import AgentConfig, AtomicAgent, BaseIOSchema, BasicChatInputSchema, BasicChatOutputSchema +from atomic_agents.agents.atomic_agent import _TOOL_MODES +from atomic_agents.context import ChatHistory + +DUMMY_KEY = "not-a-real-key" +LOGGER_NAME = "atomic_agents.agents.atomic_agent" +COUNTING_MODEL = "gpt-4o-mini" +TOOL_REQUEST_PARAMS = ("tools", "functions", "toolConfig") + + +class InputWithImage(BaseIOSchema): + """Input with image.""" + + chat_message: str = Field(..., description="Text input") + image: instructor.Image = Field(..., description="Image to analyze") + + +def _client(mode): + """Build a real Instructor client with the given mode (no request is ever made).""" + return instructor.from_openai(OpenAI(api_key=DUMMY_KEY, base_url="https://example.invalid/v1"), mode=mode) + + +def _mock_client(mode): + """A client double that exposes `mode` like a real Instructor client does.""" + client = Mock(spec=instructor.Instructor) + client.mode = mode + return client + + +def _agent(config): + return AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) + + +def _history_with_image(): + history = ChatHistory() + history.add_message( + "user", InputWithImage(chat_message="Look", image=instructor.Image.from_url("https://example.com/x.png")) + ) + return history + + +def _prepared_request(mode, messages=None): + """Return the kwargs Instructor would send for `mode`, prepared without any provider call.""" + if messages is None: + messages = [{"role": "user", "content": "hello"}] + return handle_response_model( + response_model=BasicChatOutputSchema, + mode=mode, + messages=[dict(message) for message in messages], + model="test-model", + )[1] + + +class TestModeDerivation: + """AgentConfig.mode defaults to the client's mode.""" + + def test_none_derives_client_mode(self): + agent = _agent(AgentConfig(client=_client(instructor.Mode.JSON), model=COUNTING_MODEL)) + + assert agent.mode == instructor.Mode.JSON + + def test_none_derives_tools_mode(self): + agent = _agent(AgentConfig(client=_client(instructor.Mode.TOOLS), model=COUNTING_MODEL)) + + assert agent.mode == instructor.Mode.TOOLS + + def test_none_derives_native_provider_mode(self): + agent = _agent(AgentConfig(client=_mock_client(instructor.Mode.ANTHROPIC_TOOLS), model=COUNTING_MODEL)) + + assert agent.mode == instructor.Mode.ANTHROPIC_TOOLS + + def test_derived_mode_drives_token_accounting(self): + json_agent = _agent(AgentConfig(client=_client(instructor.Mode.JSON), model=COUNTING_MODEL)) + tools_agent = _agent(AgentConfig(client=_client(instructor.Mode.TOOLS), model=COUNTING_MODEL)) + + assert json_agent._build_tools_definition() is None + assert tools_agent._build_tools_definition() is not None + + def test_derived_native_mode_counts_schema_as_tools(self): + agent = _agent(AgentConfig(client=_mock_client(instructor.Mode.ANTHROPIC_TOOLS), model=COUNTING_MODEL)) + + assert agent._build_tools_definition() is not None + assert agent.get_context_token_count().tools > 0 + + +class TestClientWithoutMode: + """Clients that expose no Mode attribute fall back to TOOLS without warnings.""" + + def test_mock_client_falls_back_to_tools_mode(self, caplog): + mock_client = Mock(spec=instructor.Instructor) # spec exposes no instance attribute `mode` + agent = _agent(AgentConfig(client=mock_client, model=COUNTING_MODEL)) + + assert agent.mode == instructor.Mode.TOOLS + assert not [r for r in caplog.records if r.name == LOGGER_NAME and r.levelno >= logging.WARNING] + + def test_explicit_mode_is_kept_for_mock_client(self, caplog): + mock_client = Mock(spec=instructor.Instructor) + agent = _agent(AgentConfig(client=mock_client, model=COUNTING_MODEL, mode=instructor.Mode.JSON)) + + assert agent.mode == instructor.Mode.JSON + assert not [r for r in caplog.records if r.name == LOGGER_NAME and r.levelno >= logging.WARNING] + + +class TestModeConsistencyWarning: + """Explicit modes are checked against how the client transmits the schema.""" + + def test_matching_mode_does_not_warn(self, caplog): + with caplog.at_level(logging.WARNING, logger=LOGGER_NAME): + agent = _agent(AgentConfig(client=_client(instructor.Mode.JSON), model=COUNTING_MODEL, mode=instructor.Mode.JSON)) + + assert agent.mode == instructor.Mode.JSON + assert not [r for r in caplog.records if r.name == LOGGER_NAME] + + def test_mismatching_accounting_warns(self, caplog): + with caplog.at_level(logging.WARNING, logger=LOGGER_NAME): + agent = _agent(AgentConfig(client=_client(instructor.Mode.JSON), model=COUNTING_MODEL, mode=instructor.Mode.TOOLS)) + + assert agent.mode == instructor.Mode.TOOLS + warnings = [r for r in caplog.records if r.name == LOGGER_NAME and r.levelno == logging.WARNING] + assert len(warnings) == 1 + assert "TOOLS" in warnings[0].getMessage() + assert "JSON" in warnings[0].getMessage() + + def test_same_accounting_different_mode_does_not_warn(self, caplog): + # TOOLS_STRICT on the config and TOOLS on the client both send a tools definition. + with caplog.at_level(logging.WARNING, logger=LOGGER_NAME): + agent = _agent( + AgentConfig(client=_client(instructor.Mode.TOOLS), model=COUNTING_MODEL, mode=instructor.Mode.TOOLS_STRICT) + ) + + assert agent.mode == instructor.Mode.TOOLS_STRICT + assert not [r for r in caplog.records if r.name == LOGGER_NAME] + + def test_native_provider_mode_matches_tools_accounting(self, caplog): + # ANTHROPIC_TOOLS on the client and TOOLS on the config are counted the same way. + with caplog.at_level(logging.WARNING, logger=LOGGER_NAME): + agent = _agent( + AgentConfig( + client=_mock_client(instructor.Mode.ANTHROPIC_TOOLS), model=COUNTING_MODEL, mode=instructor.Mode.TOOLS + ) + ) + + assert agent.mode == instructor.Mode.TOOLS + assert not [r for r in caplog.records if r.name == LOGGER_NAME] + + def test_native_provider_mode_against_json_accounting_warns(self, caplog): + with caplog.at_level(logging.WARNING, logger=LOGGER_NAME): + _agent( + AgentConfig( + client=_mock_client(instructor.Mode.ANTHROPIC_TOOLS), model=COUNTING_MODEL, mode=instructor.Mode.JSON + ) + ) + + message = [r for r in caplog.records if r.name == LOGGER_NAME][0].getMessage() + assert "AgentConfig.mode (JSON)" in message + assert "client's mode (ANTHROPIC_TOOLS)" in message + + def test_warning_mentions_both_directions(self, caplog): + with caplog.at_level(logging.WARNING, logger=LOGGER_NAME): + _agent(AgentConfig(client=_client(instructor.Mode.TOOLS), model=COUNTING_MODEL, mode=instructor.Mode.JSON)) + + message = [r for r in caplog.records if r.name == LOGGER_NAME][0].getMessage() + assert "AgentConfig.mode (JSON)" in message + assert "client's mode (TOOLS)" in message + + +class TestMultimodalTokenCounting: + """Media serialization stays in the chat format the token counter accepts.""" + + def test_image_history_counts_in_responses_mode(self): + # Regression check: Responses-format input_image parts used to reach the + # chat token counter and make get_context_token_count() raise. + agent = _agent( + AgentConfig(client=_client(instructor.Mode.RESPONSES_TOOLS), model=COUNTING_MODEL, history=_history_with_image()) + ) + + result = agent.get_context_token_count() + + assert result.history > 0 + assert result.total > 0 + + def test_media_parts_stay_chat_format_in_responses_mode(self): + agent = _agent( + AgentConfig(client=_client(instructor.Mode.RESPONSES_TOOLS), model=COUNTING_MODEL, history=_history_with_image()) + ) + + serialized = agent._serialize_history_for_token_count() + + assert [part["type"] for part in serialized[0]["content"]] == ["text", "image_url"] + + def test_image_history_counts_in_tools_mode(self): + agent = _agent(AgentConfig(client=_client(instructor.Mode.TOOLS), model=COUNTING_MODEL, history=_history_with_image())) + + result = agent.get_context_token_count() + + assert result.history > 0 + assert result.tools > 0 + + +class TestAccountingMatchesPreparedRequests: + """The schema must be counted the same way Instructor transmits it.""" + + @pytest.mark.parametrize("mode", list(instructor.Mode), ids=lambda mode: mode.name) + def test_tools_accounting_matches_prepared_request(self, mode): + try: + kwargs = _prepared_request(mode) + except Exception as e: + # Preparing some modes needs their provider SDK (not installed here) or an + # Iterable response model (parallel tool calls); those cannot be checked + # offline, and the accounting set covers them by construction. + pytest.skip(f"{mode.name} cannot be prepared without its provider setup: {type(e).__name__}: {e}") + + sends_tool_definition = any(param in kwargs for param in TOOL_REQUEST_PARAMS) + + assert sends_tool_definition == (mode in _TOOL_MODES) + + def test_anthropic_tools_request_carries_tool_without_schema_prose(self): + kwargs = _prepared_request( + instructor.Mode.ANTHROPIC_TOOLS, + messages=[{"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "hello"}], + ) + + assert kwargs["tools"] + assert kwargs["tool_choice"] + assert not any("json_schema" in str(block.get("text", "")) for block in kwargs["system"]) + assert not any("json_schema" in str(message.get("content", "")) for message in kwargs["messages"]) diff --git a/claude-plugin/atomic-agents/skills/framework/references/providers.md b/claude-plugin/atomic-agents/skills/framework/references/providers.md index 86b32789..9afef299 100644 --- a/claude-plugin/atomic-agents/skills/framework/references/providers.md +++ b/claude-plugin/atomic-agents/skills/framework/references/providers.md @@ -29,7 +29,14 @@ Model IDs in this file are illustrative. Provider model names change often — c | OpenRouter | `instructor.from_openai(OpenAI(base_url=..., api_key=...))` | `Mode.TOOLS` | `"assistant"` | OpenAI-compatible | | MiniMax | `instructor.from_openai(OpenAI(base_url=..., api_key=...), mode=Mode.JSON)` | `Mode.JSON` | `"assistant"` | OpenAI-compatible | -The `mode` value is passed to `AgentConfig(mode=...)` **and** sometimes to the Instructor factory itself. Match them. +The `mode` value is passed to `AgentConfig(mode=...)` **and** sometimes to the Instructor factory itself. These are two different knobs: + +- The **factory** mode (e.g. `instructor.from_openai(..., mode=Mode.JSON)`) decides the API call format. +- `AgentConfig(mode=...)` decides how the output schema is accounted for when counting tokens: modes that transmit the schema as a tool definition count it as a tool, the rest append it to the system message. It defaults to following the client. + +They must use the same mode family. `AgentConfig(mode=...)` is optional — omit it and the client's mode is used, so setting the factory mode alone is enough. Setting it explicitly to a mode from a different family than the client's logs a warning. + +This matters most for OpenAI-compatible hosts that do not speak the tools protocol: the factory needs `Mode.JSON`, and the accounting follows once `AgentConfig.mode` is left at its default or set to the same mode. ## OpenAI diff --git a/docs/api/agents.md b/docs/api/agents.md index d98d2ace..9cee9ed2 100644 --- a/docs/api/agents.md +++ b/docs/api/agents.md @@ -133,6 +133,7 @@ class AgentConfig: model: str = "gpt-4-turbo-preview" # Model to use history: Optional[ChatHistory] = None # History component system_prompt_generator: Optional[SystemPromptGenerator] = None # Prompt generator + mode: Optional[Mode] = None # Instructor mode for token accounting; follows the client when None input_schema: Optional[Type[BaseModel]] = None # Custom input schema output_schema: Optional[Type[BaseModel]] = None # Custom output schema model_api_parameters: Optional[dict] = None # Additional API parameters diff --git a/docs/guides/quickstart.md b/docs/guides/quickstart.md index d37ff570..2b01a960 100644 --- a/docs/guides/quickstart.md +++ b/docs/guides/quickstart.md @@ -422,6 +422,13 @@ export GEMINI_API_KEY="your-gemini-key" export OPENROUTER_API_KEY="your-openrouter-key" ``` +> **Instructor modes for OpenAI-compatible hosts:** hosts that do not speak the tools protocol — custom +> `base_url` endpoints, local servers, and gateways — need `mode=instructor.Mode.JSON` on the Instructor +> factory. `AgentConfig(mode=...)` follows the client automatically, so setting the factory mode alone is +> enough; setting it explicitly to a mode from a different family than the client's logs a warning. See the +> providers reference (`claude-plugin/atomic-agents/skills/framework/references/providers.md`) for the +> per-provider values. + ## Running the Examples To run any of these examples: From ef38ea63414eefc4f71cc03281645ee242ce5359 Mon Sep 17 00:00:00 2001 From: Andiii208 <246843242+Andiii208@users.noreply.github.com> Date: Mon, 5 Oct 2026 17:56:10 +0800 Subject: [PATCH 2/5] test(agents): cover token-count media fallback branches _serialize_history_for_token_count is a touched function, and the per-function coverage gate applies to it: it measured 75% (6 missing statements), all of them the media-serialization-failure and unknown-part fallback branches. Add two tests exercising both fallbacks, bringing the function to 100% (403 -> 405 passed). --- .../agents/test_agent_mode_consistency.py | 29 ++++++++++++++++++- 1 file changed, 28 insertions(+), 1 deletion(-) diff --git a/atomic-agents/tests/agents/test_agent_mode_consistency.py b/atomic-agents/tests/agents/test_agent_mode_consistency.py index 99168a96..9919a754 100644 --- a/atomic-agents/tests/agents/test_agent_mode_consistency.py +++ b/atomic-agents/tests/agents/test_agent_mode_consistency.py @@ -8,7 +8,7 @@ """ import logging -from unittest.mock import Mock +from unittest.mock import Mock, patch import instructor import pytest @@ -214,6 +214,33 @@ def test_image_history_counts_in_tools_mode(self): assert result.history > 0 assert result.tools > 0 + def test_media_serialization_failure_falls_back_to_placeholder(self, caplog): + agent = _agent(AgentConfig(client=_client(instructor.Mode.TOOLS), model=COUNTING_MODEL)) + history = [{"role": "user", "content": [instructor.Image.from_url("https://example.com/x.png")]}] + + with ( + caplog.at_level(logging.WARNING, logger=LOGGER_NAME), + patch.object(instructor.Image, "to_openai", side_effect=RuntimeError("serialization failed")), + patch.object(agent.history, "get_history", return_value=history), + ): + serialized = agent._serialize_history_for_token_count() + + assert serialized == [{"role": "user", "content": [{"type": "text", "text": "[image content]"}]}] + assert any("Failed to serialize" in record.getMessage() for record in caplog.records) + + def test_unknown_content_part_becomes_text(self): + agent = _agent(AgentConfig(client=_client(instructor.Mode.TOOLS), model=COUNTING_MODEL)) + + class UnknownPart: + def __str__(self): + return "unknown-part-7" + + history = [{"role": "user", "content": [UnknownPart()]}] + with patch.object(agent.history, "get_history", return_value=history): + serialized = agent._serialize_history_for_token_count() + + assert serialized == [{"role": "user", "content": [{"type": "text", "text": "unknown-part-7"}]}] + class TestAccountingMatchesPreparedRequests: """The schema must be counted the same way Instructor transmits it.""" From afddb6e270f1e2c91c331250c5c80962368f7dae Mon Sep 17 00:00:00 2001 From: Andiii208 <246843242+Andiii208@users.noreply.github.com> Date: Mon, 5 Oct 2026 22:31:18 +0800 Subject: [PATCH 3/5] refactor(agents): simplify touched init, mode and media paths The per-function complexity gate (CRAP must stay strictly below 6 for each modified function) left three touched functions over the bar: __init__ and _resolve_mode at cyclomatic complexity 6, and the media serialization inside _serialize_history_for_token_count at 8. Split each at meaningful boundaries so the orchestrating functions read as a sequence of named steps, with no behavior change: - __init__ delegates the Gemini tool-result-role default to _default_tool_result_role(); the explicit-None check is preserved. - _resolve_mode orchestrates _client_mode (mode detection), _modes_disagree (transmission comparison) and _warn_mode_mismatch (the warning), keeping derivation, override and warning behavior. - _serialize_history_for_token_count maps each content part through _token_count_content_part(), which dispatches media to _serialize_media_for_token_count() (chat-format normalization plus the placeholder fallback on serialization failure). Measured with radon (radon cc -s) on the touched file, per function, with atomic_agent.py at 100% statement coverage from the full suite: __init__ 5 -> CRAP 5 _default_tool_result_role (new) 2 _client_mode (new) 2 _modes_disagree (new) 2 _warn_mode_mismatch (new) 1 _resolve_mode 6 -> 4 _build_tools_definition 2 _serialize_history_for_token_count 8 -> 4 _token_count_content_part (new) 4 _serialize_media_for_token_count 2 Full suite: 405 passed, 17 skipped; black and flake8 clean. --- .../atomic_agents/agents/atomic_agent.py | 160 ++++++++++++------ 1 file changed, 107 insertions(+), 53 deletions(-) diff --git a/atomic-agents/atomic_agents/agents/atomic_agent.py b/atomic-agents/atomic_agents/agents/atomic_agent.py index 72f8ac6c..10cff776 100644 --- a/atomic-agents/atomic_agents/agents/atomic_agent.py +++ b/atomic-agents/atomic_agents/agents/atomic_agent.py @@ -236,12 +236,11 @@ def __init__(self, config: AgentConfig): self.system_prompt_generator = config.system_prompt_generator or SystemPromptGenerator() self.system_role = config.system_role self.assistant_role = config.assistant_role - if config.tool_result_role is not None: - self.tool_result_role = config.tool_result_role - else: - # Auto-detect: Gemini drops mid-conversation "system" messages, - # so default to "user" for Gemini backends (identified by assistant_role="model") - self.tool_result_role = "user" if config.assistant_role == "model" else "system" + self.tool_result_role = ( + config.tool_result_role + if config.tool_result_role is not None + else self._default_tool_result_role(config.assistant_role) + ) self.initial_history = self.history.copy() self.current_user_input = None self.mode = self._resolve_mode(config) @@ -252,12 +251,60 @@ def __init__(self, config: AgentConfig): self._hook_handlers: Dict[str, List[Callable]] = {} self._hooks_enabled: bool = True + @staticmethod + def _default_tool_result_role(assistant_role: str) -> str: + """ + Picks the role for tool results when AgentConfig does not override it. + + Gemini drops mid-conversation "system" messages, so Gemini backends + (identified by assistant_role="model") receive tool results as user + messages; every other backend uses "system". + """ + return "user" if assistant_role == "model" else "system" + def reset_history(self): """ Resets the history to its initial state. """ self.history = self.initial_history.copy() + @staticmethod + def _client_mode(client: Any) -> Optional[Mode]: + """ + The mode of an Instructor client, when it exposes one. + + Raw provider clients and test doubles have no comparable `mode` attribute; + None is returned for them. + """ + mode = getattr(client, "mode", None) + return mode if isinstance(mode, Mode) else None + + @staticmethod + def _modes_disagree(config_mode: Mode, client_mode: Optional[Mode]) -> bool: + """ + Whether the two modes transmit the output schema differently. + + Only the transmission matters for token accounting: two modes that both + send the schema as a tool definition are accounted for identically even + when they target different providers. + """ + if client_mode is None: + return False + return (config_mode in _TOOL_MODES) != (client_mode in _TOOL_MODES) + + @staticmethod + def _warn_mode_mismatch(config_mode: Mode, client_mode: Mode) -> None: + """Warns that token accounting will silently disagree with the API calls.""" + logging.getLogger(__name__).warning( + "AgentConfig.mode (%s) and the Instructor client's mode (%s) disagree on how the " + "output schema is transmitted, so token accounting will silently disagree with " + "the actual API calls. Create the Instructor client with %s as well, or omit " + "AgentConfig.mode to follow the client automatically.", + config_mode.name, + client_mode.name, + config_mode.name, + ) + @staticmethod def _resolve_mode(config: AgentConfig) -> Mode: """ @@ -275,26 +322,13 @@ def _resolve_mode(config: AgentConfig) -> Mode: Returns: Mode: The mode to use for token accounting. """ - client_mode = getattr(config.client, "mode", None) - if not isinstance(client_mode, Mode): - # Clients that do not expose a Mode (or test doubles of them) cannot be - # compared; the explicitly configured mode, if any, is used as-is. - client_mode = None + client_mode = AtomicAgent._client_mode(config.client) if config.mode is None: return client_mode if client_mode is not None else Mode.TOOLS - if client_mode is not None and (config.mode in _TOOL_MODES) != (client_mode in _TOOL_MODES): - logger = logging.getLogger(__name__) - logger.warning( - "AgentConfig.mode (%s) and the Instructor client's mode (%s) disagree on how the " - "output schema is transmitted, so token accounting will silently disagree with " - "the actual API calls. Create the Instructor client with %s as well, or omit " - "AgentConfig.mode to follow the client automatically.", - config.mode.name, - client_mode.name, - config.mode.name, - ) + if AtomicAgent._modes_disagree(config.mode, client_mode): + AtomicAgent._warn_mode_mismatch(config.mode, client_mode) return config.mode def add_tool_result(self, content: BaseIOSchema) -> None: @@ -524,43 +558,63 @@ def _serialize_history_for_token_count(self) -> List[Dict[str, Any]]: content = message.get("content") if isinstance(content, list): - # Multimodal content - convert to OpenAI format - serialized_content = [] - for item in content: - if isinstance(item, str): - # Text content - wrap in OpenAI text format - serialized_content.append({"type": "text", "text": item}) - elif isinstance(item, (Image, Audio, PDF)): - # Multimodal object - serialize with Instructor. The token - # counter only accepts chat-format content parts, so media is - # normalized to that format regardless of the agent's mode - # (Responses-format parts such as input_image raise). - try: - serialized_content.append(item.to_openai(Mode.JSON)) - except Exception as e: - # Log the error and use placeholder for token estimation - logger = logging.getLogger(__name__) - media_type = type(item).__name__ - logger.warning( - f"Failed to serialize {media_type} for token counting: {e}. " - f"Using placeholder for estimation." - ) - serialized_content.append({"type": "text", "text": f"[{media_type.lower()} content]"}) - elif isinstance(item, dict): - # get_history() emits pre-lowered content-part dicts (e.g. video). - # LiteLLM's token counter only accepts text and image_url parts, - # so estimate them with a placeholder. - serialized_content.append({"type": "text", "text": f"[{item.get('type', 'unknown')} content]"}) - else: - # Unknown type - convert to string - serialized_content.append({"type": "text", "text": str(item)}) - serialized.append({"role": message["role"], "content": serialized_content}) + # Multimodal content - convert to the chat format the counter accepts + serialized.append( + { + "role": message["role"], + "content": [self._token_count_content_part(item) for item in content], + } + ) else: # Simple text content - keep as is serialized.append(message) return serialized + @staticmethod + def _token_count_content_part(item: Any) -> Dict[str, Any]: + """ + Converts one history content part to the chat format the token counter accepts. + + Args: + item: A text part, an instructor multimodal object, or a content-part dict. + + Returns: + Dict[str, Any]: The content part in LiteLLM-compatible format. + """ + if isinstance(item, str): + # Text content - wrap in OpenAI text format + return {"type": "text", "text": item} + if isinstance(item, (Image, Audio, PDF)): + return AtomicAgent._serialize_media_for_token_count(item) + if isinstance(item, dict): + # get_history() emits pre-lowered content-part dicts (e.g. video). + # LiteLLM's token counter only accepts text and image_url parts, + # so estimate them with a placeholder. + return {"type": "text", "text": f"[{item.get('type', 'unknown')} content]"} + # Unknown type - convert to string + return {"type": "text", "text": str(item)} + + @staticmethod + def _serialize_media_for_token_count(item: Any) -> Dict[str, Any]: + """ + Serializes an instructor media object to the chat format the token counter accepts. + + LiteLLM's counter only accepts chat-format content parts, so media is normalized + to that format regardless of the agent's mode (Responses-format parts such as + input_image raise). Serialization failures degrade to a text placeholder so + counting still works. + """ + try: + return item.to_openai(Mode.JSON) + except Exception as e: + # Log the error and use placeholder for token estimation + media_type = type(item).__name__ + logging.getLogger(__name__).warning( + f"Failed to serialize {media_type} for token counting: {e}. Using placeholder for estimation." + ) + return {"type": "text", "text": f"[{media_type.lower()} content]"} + def get_context_token_count(self) -> TokenCountResult: """ Get the accurate token count for the current context. From 97d61df544669f9c11d8c1b8340fda3ae6a87361 Mon Sep 17 00:00:00 2001 From: Andiii208 <246843242+Andiii208@users.noreply.github.com> Date: Thu, 8 Oct 2026 22:23:32 +0800 Subject: [PATCH 4/5] refactor(agents): type the new mode and media helper contracts Replace the three new Any boundaries with accurate types: _client_mode takes the Instructor client AgentConfig actually validates, and the content-part helpers take the media union they dispatch on. --- .../atomic_agents/agents/atomic_agent.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/atomic-agents/atomic_agents/agents/atomic_agent.py b/atomic-agents/atomic_agents/agents/atomic_agent.py index 10cff776..bda921ee 100644 --- a/atomic-agents/atomic_agents/agents/atomic_agent.py +++ b/atomic-agents/atomic_agents/agents/atomic_agent.py @@ -2,7 +2,7 @@ from instructor import Mode from instructor.processing.multimodal import Image, Audio, PDF from pydantic import BaseModel, Field -from typing import Optional, Type, Generator, AsyncGenerator, get_args, get_origin, Dict, List, Callable, Any +from typing import Optional, Type, Generator, AsyncGenerator, get_args, get_origin, Dict, List, Callable, Union, Any import logging from atomic_agents.context.chat_history import ChatHistory from atomic_agents.context.base_chat_history import BaseChatHistory @@ -51,6 +51,9 @@ } ) +# Instructor multimodal objects that can be lowered to a chat-format content part. +_InstructorMedia = Union[Image, Audio, PDF] + def model_from_chunks_patched(cls, json_chunks, **kwargs): potential_object = "" @@ -269,12 +272,13 @@ def reset_history(self): self.history = self.initial_history.copy() @staticmethod - def _client_mode(client: Any) -> Optional[Mode]: + def _client_mode(client: instructor.Instructor) -> Optional[Mode]: """ The mode of an Instructor client, when it exposes one. - Raw provider clients and test doubles have no comparable `mode` attribute; - None is returned for them. + `AgentConfig` only accepts Instructor clients, but a client double or a + custom client can still lack a comparable `mode` attribute; None is + returned for them. """ mode = getattr(client, "mode", None) return mode if isinstance(mode, Mode) else None @@ -572,12 +576,13 @@ def _serialize_history_for_token_count(self) -> List[Dict[str, Any]]: return serialized @staticmethod - def _token_count_content_part(item: Any) -> Dict[str, Any]: + def _token_count_content_part(item: Union[str, _InstructorMedia, Dict[str, Any]]) -> Dict[str, Any]: """ Converts one history content part to the chat format the token counter accepts. Args: item: A text part, an instructor multimodal object, or a content-part dict. + Anything else is counted as its string form. Returns: Dict[str, Any]: The content part in LiteLLM-compatible format. @@ -596,7 +601,7 @@ def _token_count_content_part(item: Any) -> Dict[str, Any]: return {"type": "text", "text": str(item)} @staticmethod - def _serialize_media_for_token_count(item: Any) -> Dict[str, Any]: + def _serialize_media_for_token_count(item: _InstructorMedia) -> Dict[str, Any]: """ Serializes an instructor media object to the chat format the token counter accepts. From dbf421edec7a4fb739fa19ee5d63ea0e27666694 Mon Sep 17 00:00:00 2001 From: Andiii208 <246843242+Andiii208@users.noreply.github.com> Date: Thu, 8 Oct 2026 22:23:32 +0800 Subject: [PATCH 5/5] test(agents): assert mode warnings through a shared oracle test_mismatching_accounting_warns measured CRAP 8: the caplog filtering and the per-fragment assertions all lived in the test body. Extract them into _agent_warnings/_assert_warns_once so every assertion is kept and each function measures on its own, and name the parametrization id function so it can be measured too. --- .../agents/test_agent_mode_consistency.py | 35 +++++++++++++------ 1 file changed, 24 insertions(+), 11 deletions(-) diff --git a/atomic-agents/tests/agents/test_agent_mode_consistency.py b/atomic-agents/tests/agents/test_agent_mode_consistency.py index 9919a754..2dcda0e9 100644 --- a/atomic-agents/tests/agents/test_agent_mode_consistency.py +++ b/atomic-agents/tests/agents/test_agent_mode_consistency.py @@ -69,6 +69,26 @@ def _prepared_request(mode, messages=None): )[1] +def _mode_id(mode): + """Readable test id for each Instructor mode.""" + return mode.name + + +def _agent_warnings(caplog): + """Assertion oracle: the warning records the agent module emitted during a test.""" + return [record for record in caplog.records if record.name == LOGGER_NAME and record.levelno == logging.WARNING] + + +def _assert_warns_once(caplog, *fragments): + """Assertion oracle: exactly one agent warning was logged, mentioning every fragment.""" + warnings = _agent_warnings(caplog) + + assert len(warnings) == 1 + message = warnings[0].getMessage() + for fragment in fragments: + assert fragment in message + + class TestModeDerivation: """AgentConfig.mode defaults to the client's mode.""" @@ -134,10 +154,7 @@ def test_mismatching_accounting_warns(self, caplog): agent = _agent(AgentConfig(client=_client(instructor.Mode.JSON), model=COUNTING_MODEL, mode=instructor.Mode.TOOLS)) assert agent.mode == instructor.Mode.TOOLS - warnings = [r for r in caplog.records if r.name == LOGGER_NAME and r.levelno == logging.WARNING] - assert len(warnings) == 1 - assert "TOOLS" in warnings[0].getMessage() - assert "JSON" in warnings[0].getMessage() + _assert_warns_once(caplog, "TOOLS", "JSON") def test_same_accounting_different_mode_does_not_warn(self, caplog): # TOOLS_STRICT on the config and TOOLS on the client both send a tools definition. @@ -169,17 +186,13 @@ def test_native_provider_mode_against_json_accounting_warns(self, caplog): ) ) - message = [r for r in caplog.records if r.name == LOGGER_NAME][0].getMessage() - assert "AgentConfig.mode (JSON)" in message - assert "client's mode (ANTHROPIC_TOOLS)" in message + _assert_warns_once(caplog, "AgentConfig.mode (JSON)", "client's mode (ANTHROPIC_TOOLS)") def test_warning_mentions_both_directions(self, caplog): with caplog.at_level(logging.WARNING, logger=LOGGER_NAME): _agent(AgentConfig(client=_client(instructor.Mode.TOOLS), model=COUNTING_MODEL, mode=instructor.Mode.JSON)) - message = [r for r in caplog.records if r.name == LOGGER_NAME][0].getMessage() - assert "AgentConfig.mode (JSON)" in message - assert "client's mode (TOOLS)" in message + _assert_warns_once(caplog, "AgentConfig.mode (JSON)", "client's mode (TOOLS)") class TestMultimodalTokenCounting: @@ -245,7 +258,7 @@ def __str__(self): class TestAccountingMatchesPreparedRequests: """The schema must be counted the same way Instructor transmits it.""" - @pytest.mark.parametrize("mode", list(instructor.Mode), ids=lambda mode: mode.name) + @pytest.mark.parametrize("mode", list(instructor.Mode), ids=_mode_id) def test_tools_accounting_matches_prepared_request(self, mode): try: kwargs = _prepared_request(mode)