fix(cache): include tool_choice in ChatCompletionCache 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: 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,9 +193,12 @@ def _check_cache(
         elif isinstance(json_output, bool):
             json_output_data = json_output
 
+        tool_choice_data: str | ToolSchema = tool_choice.schema if isinstance(tool_choice, Tool) else tool_choice
+
         data = {
             "messages": [message.model_dump() for message in messages],
             "tools": [(tool.schema if isinstance(tool, Tool) else tool) for tool in tools],
+            "tool_choice": tool_choice_data,
             "json_output": json_output_data,
             "extra_create_args": extra_create_args,
         }
@@ -270,7 +274,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 +323,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

@@ -12,6 +12,7 @@
     SystemMessage,
     UserMessage,
 )
+from autogen_core.tools import FunctionTool
 from autogen_ext.models.cache import CHAT_CACHE_VALUE_TYPE, ChatCompletionCache
 from autogen_ext.models.replay import ReplayChatCompletionClient
 from pydantic import BaseModel
@@ -58,6 +59,46 @@ async def test_cache_basic_with_args() -> None:
     assert response2.content == responses[2]
 
 
+@pytest.mark.asyncio
+async def test_cache_tool_choice_is_part_of_cache_key_for_create_and_stream() -> None:
+    """Different tool choice policies must not reuse each other's cached responses."""
+
+    def lookup(value: str) -> str:
+        return value
+
+    tool = FunctionTool(lookup, description="Look up a value")
+    messages = [UserMessage(content="Choose a tool policy", source="user")]
+
+    replay_client = ReplayChatCompletionClient(["required response", "none response", "specific response"])
+    replay_client.set_cached_bool_value(False)
+    cached_client = ChatCompletionCache(replay_client)
+
+    required = await cached_client.create(messages, tools=[tool], tool_choice="required")
+    none = await cached_client.create(messages, tools=[tool], tool_choice="none")
+    assert required.content == "required response"
+    assert none.content == "none response"
+    assert not required.cached
+    assert not none.cached
+
+    specific = await cached_client.create(messages, tools=[tool], tool_choice=tool)
+    assert specific.content == "specific response"
+    assert not specific.cached
+
+    stream_replay_client = ReplayChatCompletionClient(["stream required", "stream none"])
+    stream_replay_client.set_cached_bool_value(False)
+    stream_cached_client = ChatCompletionCache(stream_replay_client)
+
+    streamed_required: list[Union[str, CreateResult]] = []
+    async for chunk in stream_cached_client.create_stream(messages, tools=[tool], tool_choice="required"):
+        streamed_required.append(chunk)
+    streamed_none: list[Union[str, CreateResult]] = []
+    async for chunk in stream_cached_client.create_stream(messages, tools=[tool], tool_choice="none"):
+        streamed_none.append(chunk)
+
+    assert streamed_required[-1].content == "stream required"  # type: ignore[union-attr]
+    assert streamed_none[-1].content == "stream none"  # type: ignore[union-attr]
+
+
 @pytest.mark.asyncio
 async def test_cache_structured_output_with_args() -> None:
     responses, prompts, system_prompt, _, cached_client = get_test_data(num_messages=4)