Repository: microsoft/autogen
Pull request: Add reply chat completion client
<!-- Thank you for your contribution! Please review https://microsoft.github.io/autogen/docs/Contribute before opening a pull request. -->

<!-- Please add a reviewer to the assignee section when you create a PR. If you don't have the access to it, we will shortly find a reviewer and assign them to your PR. -->

## Why are these changes needed?

<!-- Please give a short summary of the change and the problem this solves. -->
- Adds a `ReplayChatCompletionClient` client so we can mimic LLM API calls locally with predefined set of responses with some native implementation of usage logic

## Related issue number

<!-- For example: "Closes #1234" -->

Closes #3622
## Checks

- [x] I've included any doc changes needed for https://microsoft.github.io/autogen/. See https://microsoft.github.io/autogen/docs/Contribute#documentation to build and test documentation locally.
- [x] I've added tests (if relevant) corresponding to the changes introduced in this PR.
- [x] I've made sure all auto checks have passed.


Problem statement:
A replay mock model client for testing
### What feature would you like to be added?

Add a test utility class for mocking a ChatCompletionClient for testing. E.g.

```python
mock_client = ReplayChatCompletionClient(chat_completions=<a list of model client messages>)
# Use it as a regular model client, pop a message from the input list on each `create` or `create_stream` call. 
```

The model client should raise exception once the list of messages is exhausted. 

### Why is this needed?

We should avoid calling a real LLM as much as possible in unit tests for development velocity.

Hints from issue discussion:
@ekzhu 
This implementation could be as simple as wrapping existing [_mock_create](https://github.com/microsoft/autogen/blob/930e61306a2c4fe03b923e361fcfa21f32d9efb6/python/packages/autogen-ext/tests/models/test_openai_model_client.py#L108) and [_mock_create_stream](_mock_create_stream) into the `ReplayChatCompletionClient` all the way to (try to) calculate and return usages and support Tools.  What scope do you think we should take for this ticket so I can create subtasks accordingly.
It is sufficient to have a mock client we can reuse for each instance of model client. Currently the mock approach is for all clients rather than an instance.

Commit message: update the docstring for reply chat completion client

Signed-off-by: Mohammad Mazraeh <mazraeh.mohammad@gmail.com>

Patch to review:
diff --git a/python/packages/autogen-ext/src/autogen_ext/models/__init__.py b/python/packages/autogen-ext/src/autogen_ext/models/__init__.py
index e7b2b76ae362..d39c1d9bf247 100644
--- a/python/packages/autogen-ext/src/autogen_ext/models/__init__.py
+++ b/python/packages/autogen-ext/src/autogen_ext/models/__init__.py
@@ -2,8 +2,6 @@
     AzureOpenAIChatCompletionClient,
     OpenAIChatCompletionClient,
 )
+from ._reply_chat_completion_client import ReplayChatCompletionClient
 
-__all__ = [
-    "AzureOpenAIChatCompletionClient",
-    "OpenAIChatCompletionClient",
-]
+__all__ = ["AzureOpenAIChatCompletionClient", "OpenAIChatCompletionClient", "ReplayChatCompletionClient"]
diff --git a/python/packages/autogen-ext/src/autogen_ext/models/_reply_chat_completion_client.py b/python/packages/autogen-ext/src/autogen_ext/models/_reply_chat_completion_client.py
new file mode 100644
index 000000000000..084ccd5f7d75
--- /dev/null
+++ b/python/packages/autogen-ext/src/autogen_ext/models/_reply_chat_completion_client.py
@@ -0,0 +1,162 @@
+from __future__ import annotations
+
+import logging
+from typing import Any, AsyncGenerator, List, Mapping, Optional, Sequence, Union
+
+from autogen_core.application.logging import EVENT_LOGGER_NAME
+from autogen_core.base import CancellationToken
+from autogen_core.components.models import (
+    ChatCompletionClient,
+    CreateResult,
+    LLMMessage,
+    ModelCapabilities,
+    RequestUsage,
+)
+from autogen_core.components.tools import Tool, ToolSchema
+
+logger = logging.getLogger(EVENT_LOGGER_NAME)
+
+
+class ReplayChatCompletionClient:
+    """A mock chat completion client that replays predefined responses.
+
+    This class simulates a chat completion client by replaying a predefined list of responses. It supports both single completion and streaming responses. The responses can be either strings or CreateResult objects.
+
+    Example usage:
+    ```
+    chat_completions = [
+        "Hello, how can I assist you today?",
+        "I'm happy to help with any questions you have.",
+        "Is there anything else I can assist you with?",
+    ]
+    client = ReplayChatCompletionClient(chat_completions)
+    messages = [LLMMessage(content="What can you do?")]
+    response = await client.create(messages)
+    print(response.content)  # Output: "Hello, how can I assist you today?"
+
+    async for token in client.create_stream(messages):
+        print(token, end="")  # Output: "Hello, how can I assist you today?"
+    ```
+    """
+
+    __protocol__: ChatCompletionClient
+
+    # TODO: Support FunctionCall in responses
+    # TODO: Support logprobs in Responses
+    # TODO: Support model capabilities
+
+    def __init__(
+        self,
+        chat_completions: Sequence[Union[str, CreateResult]],
+    ):
+        self.chat_completions = list(chat_completions)
+        self._cur_usage = RequestUsage(prompt_tokens=0, completion_tokens=0)
+        self._total_usage = RequestUsage(prompt_tokens=0, completion_tokens=0)
+        self.provided_message_count = len(self.chat_completions)
+        self._model_capabilities = ModelCapabilities(vision=False, function_calling=False, json_output=False)
+        self._total_available_tokens = 10000
+
+    async def create(
+        self,
+        messages: Sequence[LLMMessage],
+        tools: Sequence[Tool | ToolSchema] = [],
+        json_output: Optional[bool] = None,
+        extra_create_args: Mapping[str, Any] = {},
+        cancellation_token: Optional[CancellationToken] = None,
+    ) -> CreateResult:
+        """Return the next completion from the list."""
+        if not self.chat_completions:
+            raise ValueError("No more mock responses available")
+
+        response = self.chat_completions.pop(0)
+        _, prompt_token_count = self._tokenize(messages)
+        if isinstance(response, str):
+            _, output_token_count = self._tokenize(response)
+            self._cur_usage = RequestUsage(prompt_tokens=prompt_token_count, completion_tokens=output_token_count)
+            response = CreateResult(finish_reason="stop", content=response, usage=self._cur_usage, cached=True)
+        else:
+            self._cur_usage = RequestUsage(
+                prompt_tokens=prompt_token_count, completion_tokens=response.usage.completion_tokens
+            )
+
+        self._update_total_usage()
+        return response
+
+    async def create_stream(
+        self,
+        messages: Sequence[LLMMessage],
+        tools: Sequence[Tool | ToolSchema] = [],
+        json_output: Optional[bool] = None,
+        extra_create_args: Mapping[str, Any] = {},
+        cancellation_token: Optional[CancellationToken] = None,
+    ) -> AsyncGenerator[Union[str, CreateResult], None]:
+        """Return the next completion as a stream."""
+        if not self.chat_completions:
+            raise ValueError("No more mock responses available")
+
+        response = self.chat_completions.pop(0)
+        _, prompt_token_count = self._tokenize(messages)
+        if isinstance(response, str):
+            output_tokens, output_token_count = self._tokenize(response)
+            self._cur_usage = RequestUsage(prompt_tokens=prompt_token_count, completion_tokens=output_token_count)
+
+            for i, token in enumerate(output_tokens):
+                if i < len(output_tokens) - 1:
+                    yield token + " "
+                else:
+                    yield token
+            self._update_total_usage()
+        else:
+            self._cur_usage = RequestUsage(
+                prompt_tokens=prompt_token_count, completion_tokens=response.usage.completion_tokens
+            )
+            yield response
+            self._update_total_usage()
+
+    def actual_usage(self) -> RequestUsage:
+        return self._cur_usage
+
+    def total_usage(self) -> RequestUsage:
+        return self._total_usage
+
+    def count_tokens(self, messages: Sequence[LLMMessage], tools: Sequence[Tool | ToolSchema] = []) -> int:
+        _, token_count = self._tokenize(messages)
+        return token_count
+
+    def remaining_tokens(self, messages: Sequence[LLMMessage], tools: Sequence[Tool | ToolSchema] = []) -> int:
+        return max(
+            0, self._total_available_tokens - self._total_usage.prompt_tokens - self._total_usage.completion_tokens
+        )
+
+    def _tokenize(self, messages: Union[str, LLMMessage, Sequence[LLMMessage]]) -> tuple[list[str], int]:
+        total_tokens = 0
+        all_tokens: List[str] = []
+        if isinstance(messages, str):
+            tokens = messages.split()
+            total_tokens += len(tokens)
+            all_tokens.extend(tokens)
+        elif hasattr(messages, "content"):
+            if isinstance(messages.content, str):  # type: ignore [reportAttributeAccessIssue]
+                tokens = messages.content.split()  # type: ignore [reportAttributeAccessIssue]
+                total_tokens += len(tokens)
+                all_tokens.extend(tokens)
+            else:
+                logger.warning("Token count has been done only on string content", RuntimeWarning)
+        elif isinstance(messages, Sequence):
+            for message in messages:
+                if isinstance(message.content, str):  # type: ignore [reportAttributeAccessIssue, union-attr]
+                    tokens = message.content.split()  # type: ignore [reportAttributeAccessIssue, union-attr]
+                    total_tokens += len(tokens)
+                    all_tokens.extend(tokens)
+                else:
+                    logger.warning("Token count has been done only on string content", RuntimeWarning)
+        return all_tokens, total_tokens
+
+    def _update_total_usage(self) -> None:
+        self._total_usage.completion_tokens += self._cur_usage.completion_tokens
+        self._total_usage.prompt_tokens += self._cur_usage.prompt_tokens
+
+    @property
+    def capabilities(self) -> ModelCapabilities:
+        """Return mock capabilities."""
+        return self._model_capabilities
diff --git a/python/packages/autogen-ext/tests/models/test_reply_chat_completion_client.py b/python/packages/autogen-ext/tests/models/test_reply_chat_completion_client.py
new file mode 100644
index 000000000000..f884faabf231
--- /dev/null
+++ b/python/packages/autogen-ext/tests/models/test_reply_chat_completion_client.py
@@ -0,0 +1,148 @@
+import copy
+from dataclasses import dataclass
+from typing import List
+
+import pytest
+from autogen_core.application import SingleThreadedAgentRuntime
+from autogen_core.base import AgentId, MessageContext
+from autogen_core.components import DefaultTopicId, RoutedAgent, default_subscription, message_handler
+from autogen_core.components.models import ChatCompletionClient, CreateResult, SystemMessage, UserMessage
+from autogen_ext.models import ReplayChatCompletionClient
+
+
+@dataclass
+class ContentMessage:
+    content: str
+
+
+class LLMAgent(RoutedAgent):
+    def __init__(self, model_client: ChatCompletionClient) -> None:
+        super().__init__("LLM Agent!")
+        self._chat_history: List[ContentMessage] = []
+        self._model_client = model_client
+        self.num_calls = 0
+
+    @message_handler
+    async def on_new_message(self, message: ContentMessage, ctx: MessageContext) -> None:
+        self._chat_history.append(message)
+        self.num_calls += 1
+        completion = await self._model_client.create(messages=self._fixed_message_history_type)
+        if isinstance(completion.content, str):
+            await self.publish_message(ContentMessage(content=completion.content), DefaultTopicId())
+        else:
+            raise TypeError(f"Completion content of type {type(completion.content)} is not supported")
+
+    @property
+    def _fixed_message_history_type(self) -> List[SystemMessage]:
+        return [SystemMessage(msg.content) for msg in self._chat_history]
+
+
+@default_subscription
+class LLMAgentWithDefaultSubscription(LLMAgent): ...
+
+
+@pytest.mark.asyncio
+async def test_reply_chat_completion_client() -> None:
+    num_messages = 5
+    messages = [f"Message {i}" for i in range(num_messages)]
+    reply_model_client = ReplayChatCompletionClient(messages)
+
+    for i in range(num_messages):
+        completion: CreateResult = await reply_model_client.create([UserMessage(content="dummy", source="_")])
+        assert completion.content == messages[i]
+    with pytest.raises(ValueError, match="No more mock responses available"):
+        await reply_model_client.create([UserMessage(content="dummy", source="_")])
+
+
+@pytest.mark.asyncio
+async def test_reply_chat_completion_client_create_stream() -> None:
+    num_messages = 5
+    messages = [f"Message {i}" for i in range(num_messages)]
+    reply_model_client = ReplayChatCompletionClient(messages)
+
+    for i in range(num_messages):
+        result: List[str] = []
+        async for completion in reply_model_client.create_stream([UserMessage(content="dummy", source="_")]):
+            text = completion.content if isinstance(completion, CreateResult) else completion
+            assert isinstance(text, str)
+            result.append(text)
+        assert "".join(result) == messages[i]
+
+    with pytest.raises(ValueError, match="No more mock responses available"):
+        await reply_model_client.create([UserMessage(content="dummy", source="_")])
+
+
+@pytest.mark.asyncio
+async def test_register_receives_publish_llm() -> None:
+    runtime = SingleThreadedAgentRuntime()
+    runtime.start()
+
+    reply_model_client_1 = ReplayChatCompletionClient(["Hi!", "Doing Good, you?", "Bye!"])
+    reply_model_client_2 = ReplayChatCompletionClient(["Hi! How are you doing?", "Good, nice to meet you, bye!"])
+
+    # First registered models gets the first message
+    assert reply_model_client_1.provided_message_count == 1 + reply_model_client_2.provided_message_count
+
+    await LLMAgentWithDefaultSubscription.register(
+        runtime, "LLMAgent1", lambda: LLMAgentWithDefaultSubscription(reply_model_client_1)
+    )
+
+    await LLMAgentWithDefaultSubscription.register(
+        runtime, "LLMAgent2", lambda: LLMAgentWithDefaultSubscription(reply_model_client_2)
+    )
+
+    await runtime.publish_message(ContentMessage(content="Let's get started!"), DefaultTopicId())
+    await runtime.stop_when_idle()
+
+    agent_1 = await runtime.try_get_underlying_agent_instance(
+        AgentId("LLMAgent1", key="default"), type=LLMAgentWithDefaultSubscription
+    )
+
+    agent_2 = await runtime.try_get_underlying_agent_instance(
+        AgentId("LLMAgent2", key="default"), type=LLMAgentWithDefaultSubscription
+    )
+
+    assert agent_1.num_calls == 1 + reply_model_client_2.provided_message_count
+    assert agent_2.num_calls == 1 + reply_model_client_1.provided_message_count
+
+
+@pytest.mark.asyncio
+async def test_token_count_logics() -> None:
+    phrases = [
+        "This is a test message.",
+        "This is another test message.",
+        "This is yet another test message.",
+        "Maybe even more messages?",
+    ]
+    reply_model_client = ReplayChatCompletionClient(phrases)
+
+    messages = [UserMessage(content="How many tokens are in this message?", source="_")]
+
+    token_count = reply_model_client.count_tokens(messages)
+    assert token_count == 7
+
+    _ = await reply_model_client.create(messages)
+    remaining_tokens = reply_model_client.remaining_tokens(messages)
+    assert remaining_tokens == 9988
+
+    multiple_messages = [UserMessage(content="This is another test message.", source="_")]
+    total_token_count = reply_model_client.count_tokens(messages + multiple_messages)
+    assert total_token_count == 12
+
+    before_cteate_usage = copy.deepcopy(reply_model_client.total_usage())
+    completion: CreateResult = await reply_model_client.create(messages)
+
+    assert completion.usage.prompt_tokens == 7
+    assert completion.usage.completion_tokens == 5
+
+    after_create_usage = reply_model_client.total_usage()
+    assert after_create_usage.prompt_tokens > before_cteate_usage.prompt_tokens
+    assert after_create_usage.completion_tokens > before_cteate_usage.completion_tokens
+
+    before_cteate_stream_usage = copy.deepcopy(reply_model_client.total_usage())
+
+    async for _ in reply_model_client.create_stream(messages):
+        pass
+    after_create_stream_usage = reply_model_client.total_usage()
+    assert after_create_stream_usage.completion_tokens > before_cteate_stream_usage.completion_tokens
+    assert after_create_stream_usage.prompt_tokens > before_cteate_stream_usage.prompt_tokens


Code review comments:
diff_hunk: @@ -0,0 +1,162 @@
+from __future__ import annotations
+
+import logging
+from typing import Any, AsyncGenerator, List, Mapping, Optional, Sequence, Union
+
+from autogen_core.application.logging import EVENT_LOGGER_NAME
+from autogen_core.base import CancellationToken
+from autogen_core.components.models import (
+    ChatCompletionClient,
+    CreateResult,
+    LLMMessage,
+    ModelCapabilities,
+    RequestUsage,
+)
+from autogen_core.components.tools import Tool, ToolSchema
+
+logger = logging.getLogger(EVENT_LOGGER_NAME)
+
+
+class ReplayChatCompletionClient:
+    """A mock chat completion client that replays predefined responses.
+
+    This class simulates a chat completion client by replaying a predefined list of responses. It supports both single completion and streaming responses. The responses can be either strings or CreateResult objects.
+
+    Example usage:
+    ```
+    chat_completions = [
+        "Hello, how can I assist you today?",
+        "I'm happy to help with any questions you have.",
+        "Is there anything else I can assist you with?",
+    ]
+    client = ReplayChatCompletionClient(chat_completions)
+    messages = [LLMMessage(content="What can you do?")]
+    response = await client.create(messages)
+    print(response.content)  # Output: "Hello, how can I assist you today?"
+
+    async for token in client.create_stream(messages):
+        print(token, end="")  # Output: "Hello, how can I assist you today?"
+    ```
+    """
+
+    __protocol__: ChatCompletionClient
+
+    # TODO: Support FunctionCall in responses
+    # TODO: Support logprobs in Responses
+    # TODO: Support model capabilities
+
+    def __init__(
+        self,
+        chat_completions: Sequence[Union[str, CreateResult]],
+    ):
+        self.chat_completions = list(chat_completions)
+        self._cur_usage = RequestUsage(prompt_tokens=0, completion_tokens=0)
+        self._total_usage = RequestUsage(prompt_tokens=0, completion_tokens=0)
+        self.provided_message_count = len(self.chat_completions)
+        self._model_capabilities = ModelCapabilities(vision=False, function_calling=False, json_output=False)
+        self._total_available_tokens = 10000
+
+    async def create(
+        self,
+        messages: Sequence[LLMMessage],
+        tools: Sequence[Tool | ToolSchema] = [],
+        json_output: Optional[bool] = None,
+        extra_create_args: Mapping[str, Any] = {},
+        cancellation_token: Optional[CancellationToken] = None,
+    ) -> CreateResult:
+        """Return the next completion from the list."""
+        if not self.chat_completions:
+            raise ValueError("No more mock responses available")
+
+        response = self.chat_completions.pop(0)
path: python/packages/autogen-ext/src/autogen_ext/models/_reply_chat_completion_client.py
text: @user1:
In many scenarios we want to reset the clients in test to get the same sequence again. So instead we can use index rather than popping the response. 

Also what happens if the responses are exhausted? Needs to raise an informative error here.

@author:
1. Makes sense, in the indexing approach the index needs to be reseted too, do you think having a `.reset` approach is good?
2. In case of messages being exhausted, currently a `ValueError("No more mock responses available")` is being thrown. See the [test of this functionality](https://github.com/microsoft/autogen/pull/4096/files#r1835754287)

@user1:
> Makes sense, in the indexing approach the index needs to be reseted too, do you think having a .reset approach is good?

Yes, I think a reset approach makes it easier to work with.