From f02a7a527d164b0d319e19943e71eb6dc4620d01 Mon Sep 17 00:00:00 2001 From: HumphreySun98 Date: Wed, 8 Jul 2026 15:12:09 -0500 Subject: [PATCH] fix: include tool_choice in ChatCompletionCache cache key `ChatCompletionCache` forwards `tool_choice` to the underlying client, where it materially changes the model's output (`"auto"` vs `"required"` vs `"none"` vs a specific forced tool). However `_check_cache` did not include `tool_choice` when computing the cache key, so two otherwise-identical calls that differed only in `tool_choice` collided and the second call silently returned the first call's cached result. Include `tool_choice` in the cache-key data (representing a forced `Tool` by its name) and pass it through from both `create` and `create_stream`. Co-Authored-By: Claude Opus 4.8 --- .../models/cache/_chat_completion_cache.py | 9 +++++- .../models/test_chat_completion_cache.py | 31 +++++++++++++++++++ 2 files changed, 39 insertions(+), 1 deletion(-) diff --git a/python/packages/autogen-ext/src/autogen_ext/models/cache/_chat_completion_cache.py b/python/packages/autogen-ext/src/autogen_ext/models/cache/_chat_completion_cache.py index 124a3ea7643c..95a6cbaef002 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/cache/_chat_completion_cache.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/cache/_chat_completion_cache.py @@ -179,6 +179,7 @@ def _check_cache( tools: Sequence[Tool | ToolSchema], json_output: Optional[bool | type[BaseModel]], extra_create_args: Mapping[str, Any], + tool_choice: Tool | Literal["auto", "required", "none"] = "auto", ) -> tuple[Optional[Union[CreateResult, List[Union[str, CreateResult]]]], str]: """ Helper function to check the cache for a result. @@ -192,11 +193,16 @@ def _check_cache( elif isinstance(json_output, bool): json_output_data = json_output + # tool_choice is forwarded to the underlying client and changes the model output, + # so it must be part of the cache key. Represent a forced Tool by its name. + tool_choice_data = tool_choice if isinstance(tool_choice, str) else tool_choice.schema["name"] + data = { "messages": [message.model_dump() for message in messages], "tools": [(tool.schema if isinstance(tool, Tool) else tool) for tool in tools], "json_output": json_output_data, "extra_create_args": extra_create_args, + "tool_choice": tool_choice_data, } serialized_data = json.dumps(data, sort_keys=True) cache_key = hashlib.sha256(serialized_data.encode()).hexdigest() @@ -270,7 +276,7 @@ async def create( NOTE: cancellation_token is ignored for cached results. """ - cached_result, cache_key = self._check_cache(messages, tools, json_output, extra_create_args) + cached_result, cache_key = self._check_cache(messages, tools, json_output, extra_create_args, tool_choice) if cached_result is not None: if isinstance(cached_result, CreateResult): # Cache hit from previous non-streaming call @@ -319,6 +325,7 @@ async def _generator() -> AsyncGenerator[Union[str, CreateResult], None]: tools, json_output, extra_create_args, + tool_choice, ) if cached_result is not None: if isinstance(cached_result, list): diff --git a/python/packages/autogen-ext/tests/models/test_chat_completion_cache.py b/python/packages/autogen-ext/tests/models/test_chat_completion_cache.py index 8627cb9f6221..9863c44776c4 100644 --- a/python/packages/autogen-ext/tests/models/test_chat_completion_cache.py +++ b/python/packages/autogen-ext/tests/models/test_chat_completion_cache.py @@ -956,3 +956,34 @@ async def test_create_stream_with_cached_non_streaming_result_non_string_content assert stream_results[0].content == [expected_function_call] assert stream_results[0].finish_reason == "function_calls" assert stream_results[0].cached is True + + +def test_check_cache_key_includes_tool_choice() -> None: + """Different ``tool_choice`` values must produce different cache keys. + + Regression test: ``tool_choice`` is forwarded to the underlying client and changes the + model output, but it was omitted from the cache key, so calls that differed only in + ``tool_choice`` collided and returned each other's cached results. + """ + from autogen_core.tools import FunctionTool + + def _dummy_tool(value: str) -> str: + return value + + tool = FunctionTool(_dummy_tool, description="dummy", name="my_tool") + + replay_client = ReplayChatCompletionClient(["response"]) + cached_client = ChatCompletionCache(replay_client) + messages: List[LLMMessage] = [UserMessage(content="hello", source="user")] + + _, key_auto = cached_client._check_cache(messages, [tool], None, {}, "auto") # type: ignore[reportPrivateUsage] + _, key_required = cached_client._check_cache(messages, [tool], None, {}, "required") # type: ignore[reportPrivateUsage] + _, key_none = cached_client._check_cache(messages, [tool], None, {}, "none") # type: ignore[reportPrivateUsage] + _, key_tool = cached_client._check_cache(messages, [tool], None, {}, tool) # type: ignore[reportPrivateUsage] + + # All four tool_choice values must yield distinct cache keys. + assert len({key_auto, key_required, key_none, key_tool}) == 4 + + # The default (no tool_choice passed) matches the explicit "auto" default. + _, key_default = cached_client._check_cache(messages, [tool], None, {}) # type: ignore[reportPrivateUsage] + assert key_default == key_auto