fix: include tool_choice in ChatCompletionCache cache key

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: Optional[Tool | Literal["auto", "required", "none"]] = None,
     ) -> tuple[Optional[Union[CreateResult, List[Union[str, CreateResult]]]], str]:
         """
         Helper function to check the cache for a result.
@@ -192,11 +193,20 @@ def _check_cache(
         elif isinstance(json_output, bool):
             json_output_data = json_output
 
+        # Include tool_choice in cache key to prevent stale cache hits
+        tool_choice_data: str | None = None
+        if tool_choice is not None:
+            if isinstance(tool_choice, str):
+                tool_choice_data = tool_choice
+            elif isinstance(tool_choice, Tool):
+                tool_choice_data = json.dumps(tool_choice.schema, sort_keys=True)
+
         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 +280,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 +329,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):

python/packages/autogen-ext/tests/models/test_chat_completion_cache.py

@@ -956,3 +956,27 @@ 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
+
+
+@pytest.mark.asyncio
+async def test_cache_key_includes_tool_choice() -> None:
+    """Test that different tool_choice values produce different cache keys."""
+    from autogen_core import InMemoryStore
+    from autogen_core.models import UserMessage
+    from autogen_ext.models.cache import ChatCompletionCache
+    from autogen_ext.models.replay import ReplayChatCompletionClient
+
+    responses = ["response 0", "response 1"]
+    replay_client = ReplayChatCompletionClient(responses)
+    replay_client.set_cached_bool_value(False)
+    cached_client = ChatCompletionCache(replay_client, InMemoryStore())
+
+    message = [UserMessage(content="hi", source="user")]
+
+    r0 = await cached_client.create(message, tool_choice="auto")
+    assert not r0.cached
+    assert r0.content == "response 0"
+
+    r1 = await cached_client.create(message, tool_choice="none")
+    assert not r1.cached  # Should be cache miss due to different tool_choice
+    assert r1.content == "response 1"