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/7] 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/7] 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/7] 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/7] 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/7] 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) From 6b6a7db33b4ed8d38f59879038a7fdf4bf71428b Mon Sep 17 00:00:00 2001 From: Andiii208 <246843242+Andiii208@users.noreply.github.com> Date: Mon, 19 Oct 2026 16:25:51 +0800 Subject: [PATCH 6/7] test(agents): cover the xAI tool representation in accounting tests Instructor 1.17 prepares xAI tools requests with the schema under the private _xai_tool kwarg instead of an OpenAI-style tools parameter, so the prepared-request accounting check did not recognize the tool and XAI_TOOLS disagreed with the accounting set. Recognize the kwarg and assert the representation carries the output schema. --- .../tests/agents/test_agent_mode_consistency.py | 13 ++++++++++++- 1 file changed, 12 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 2dcda0e9..4387294e 100644 --- a/atomic-agents/tests/agents/test_agent_mode_consistency.py +++ b/atomic-agents/tests/agents/test_agent_mode_consistency.py @@ -23,7 +23,9 @@ DUMMY_KEY = "not-a-real-key" LOGGER_NAME = "atomic_agents.agents.atomic_agent" COUNTING_MODEL = "gpt-4o-mini" -TOOL_REQUEST_PARAMS = ("tools", "functions", "toolConfig") +# Request parameters Instructor uses to transmit the output schema as a tool or +# function definition; xAI carries it under the private _xai_tool kwarg instead. +TOOL_REQUEST_PARAMS = ("tools", "functions", "toolConfig", "_xai_tool") class InputWithImage(BaseIOSchema): @@ -282,3 +284,12 @@ def test_anthropic_tools_request_carries_tool_without_schema_prose(self): 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"]) + + def test_xai_tools_request_carries_schema_as_tool(self): + # xAI has no OpenAI-style tools parameter: the schema travels as a tool + # definition under Instructor's private _xai_tool kwarg, which the + # accounting set must recognize. + kwargs = _prepared_request(instructor.Mode.XAI_TOOLS) + + assert kwargs["_xai_tool"]["name"] == BasicChatOutputSchema.__name__ + assert "chat_message" in kwargs["_xai_tool"]["parameters"]["properties"] From ec65ff22672e476b56a119f0e59c7c615e2d6f23 Mon Sep 17 00:00:00 2001 From: Andiii208 <246843242+Andiii208@users.noreply.github.com> Date: Mon, 19 Oct 2026 16:25:56 +0800 Subject: [PATCH 7/7] refactor(agents): type the token-count content-part input as object _token_count_content_part deliberately converts unknown history content parts to text, but its Union annotation excluded exactly those objects. Take the input as object and keep the isinstance narrowing, so the contract includes the supported fallback path. --- atomic-agents/atomic_agents/agents/atomic_agent.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/atomic-agents/atomic_agents/agents/atomic_agent.py b/atomic-agents/atomic_agents/agents/atomic_agent.py index 1e536b83..018c5dc0 100644 --- a/atomic-agents/atomic_agents/agents/atomic_agent.py +++ b/atomic-agents/atomic_agents/agents/atomic_agent.py @@ -581,7 +581,7 @@ def _serialize_history_for_token_count(self) -> List[Dict[str, Any]]: return serialized @staticmethod - def _token_count_content_part(item: Union[str, _InstructorMedia, Dict[str, Any]]) -> Dict[str, Any]: + def _token_count_content_part(item: object) -> Dict[str, Any]: """ Converts one history content part to the chat format the token counter accepts.