Repository: microsoft/autogen
Pull request: Feature/agentchat message id field 6317
## Why are these changes needed?

This PR implements unique ID fields for AgentChat messages to enable proper correlation between streaming chunks and completed messages. Currently, there's no way to correlate `ModelClientStreamingChunkEvent` chunks with their eventual completed message, which can lead to duplicate message display in streaming scenarios.

The implementation adds:
- `id: str` field to `BaseChatMessage` with automatic UUID generation
- `id: str` field to `BaseAgentEvent` with automatic UUID generation  
- `full_message_id: str | None` field to `ModelClientStreamingChunkEvent` for chunk-to-message correlation

This allows consumers of the streaming API to avoid double-printing messages by correlating chunks with their final complete message.

## Related issue number

Closes #6317

## Checks

- [x] I've included any doc changes needed for <https://microsoft.github.io/autogen/>. See <https://github.com/microsoft/autogen/blob/main/CONTRIBUTING.md> 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:
AgentChat messages should have an `id` field.
### Confirmation

- [x] I confirm that I am a maintainer and so can use this template. If I am not, I understand this issue will be closed and I will be asked to use a different template.

### Issue body

Message id is useful to correlate `ModelClientStreamingChunkEvent` with the eventual completed message being produced. 

This should be added to `BaseChatMessage` and `BaseAgentEvent` as an `id: str` field. For `ModelClientStreamingChunkEvent`, it should have a separate optional field `full_message_id: str | None = None` to reference the complete message that may come after the chunks. So the consumer of the stream don't need to double-print it.

Hints from issue discussion:
Hi @ekzhu, I'm interested in working on this, my initial thought was https://github.com/microsoft/autogen/pull/6351.

Currently I'm not sure if the "consumer" refers to a `Console()` call, could you help to explain a bit more? If it's, I haven't figured out the necessity of maintaining a `full_message_id`. For example, in the `Console()` call, if all `ModelClientStreamingChunkEvents` point to the same `full_message_id`, it seems like we could simply skip `message.chat_message` of the `Response` type's message field rather than comparing `message.chat_message.id` with `full_message_id` first . I think I may miss some cases where it's necessary

@thorseraq  thanks. To answer your question,

>  the "consumer" refers to a Console() call,

Yes, currently, the only consumer of `run_stream` and `on_messages_streams` are `Console` and `RichConsole` in `magentic-one-cli` package. 

> I haven't figured out the necessity of maintaining a full_message_id

So from a consumer perspective, they may skip the accumulated message if the message id matches the `full_message_id` of the chunk events. Currently, the consumer __assumes__ the full message always follows the sequence of chunk message. This may not be the case in the future, as the messages may arrive out of order when we introduce concurrency to teams. 

So introducing message id and drawing reference from chunk event to the eventual full message helps provide clarity to the consumer how to correlate messages.

Commit message: fix: used default factory unless for streaming chunk events

Patch to review:
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py
index 905df4f9d3e9..ed960ae9060d 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py
@@ -1,20 +1,11 @@
+from __future__ import annotations
+
 import asyncio
 import json
 import logging
+import uuid
 import warnings
-from typing import (
-    Any,
-    AsyncGenerator,
-    Awaitable,
-    Callable,
-    Dict,
-    List,
-    Mapping,
-    Optional,
-    Sequence,
-    Tuple,
-    Union,
-)
+from typing import Any, AsyncGenerator, Awaitable, Callable, Dict, List, Mapping, Optional, Sequence, Tuple, Union
 
 from autogen_core import CancellationToken, Component, ComponentModel, FunctionCall
 from autogen_core.memory import Memory
@@ -140,12 +131,12 @@ class AssistantAgent(BaseChatAgent, Component[AssistantAgentConfig]):
     .. tip::
 
         By default, the tool call results are returned as the response when tool
-        calls are made, so pay close attention to how the tools’ return values
+        calls are made, so pay close attention to how the tools' return values
         are formatted—especially if another agent expects a specific schema.
 
         * Use **`tool_call_summary_format`** for a simple static template.
         * Use **`tool_call_summary_formatter`** for full programmatic control
-          (e.g., “hide large success payloads, show full details on error”).
+          (e.g., "hide large success payloads, show full details on error").
 
         *Note*: `tool_call_summary_formatter` is **not serializable** and will
         be ignored when an agent is loaded from, or exported to, YAML/JSON
@@ -555,7 +546,7 @@ async def main() -> None:
             - Start with a pizza crust (store-bought or homemade).
             - Spread a layer of marinara or tomato sauce evenly over the crust.
             - Top with your favorite vegetables like bell peppers, mushrooms, onions, olives, and spinach.
-            - Add some protein if you’d like, such as grilled chicken or pepperoni (ensure it's cheese-free).
+            - Add some protein if you'd like, such as grilled chicken or pepperoni (ensure it's cheese-free).
             - Sprinkle with herbs like oregano and basil, and maybe a drizzle of olive oil.
             - Bake according to the crust instructions until the edges are golden and the veggies are cooked.
 
@@ -655,6 +646,9 @@ async def run_reasoning_agent() -> None:
     component_config_schema = AssistantAgentConfig
     component_provider_override = "autogen_agentchat.agents.AssistantAgent"
 
+    # Store correlation between CreateResult objects and message IDs
+    _message_id_map: Dict[int, str] = {}
+
     def __init__(
         self,
         name: str,
@@ -962,6 +956,12 @@ async def _call_llm(
 
         if model_client_stream:
             model_result: Optional[CreateResult] = None
+            # Generate a message ID that will be used for both chunks and the final message
+            full_message_id = str(uuid.uuid4())
+
+            # Store the ID in the model_result for later correlation
+            message_id_for_correlation = full_message_id
+
             async for chunk in model_client.create_stream(
                 llm_messages,
                 tools=tools,
@@ -970,8 +970,12 @@ async def _call_llm(
             ):
                 if isinstance(chunk, CreateResult):
                     model_result = chunk
+                    # Store the ID in the correlation map using object ID as key
+                    cls._message_id_map[id(model_result)] = message_id_for_correlation
                 elif isinstance(chunk, str):
-                    yield ModelClientStreamingChunkEvent(content=chunk, source=agent_name)
+                    yield ModelClientStreamingChunkEvent(
+                        content=chunk, source=agent_name, full_message_id=full_message_id
+                    )
                 else:
                     raise RuntimeError(f"Invalid chunk type: {type(chunk)}")
             if model_result is None:
@@ -1013,6 +1017,9 @@ async def _process_model_result(
 
         # If direct text response (string)
         if isinstance(model_result.content, str):
+            # Use the same ID for the final message as was used for streaming chunks
+            # Get the message ID from the correlation map using object ID as key
+            message_id = cls._message_id_map.get(id(model_result)) or str(uuid.uuid4())
             if output_content_type:
                 content = output_content_type.model_validate_json(model_result.content)
                 yield Response(
@@ -1021,6 +1028,7 @@ async def _process_model_result(
                         source=agent_name,
                         models_usage=model_result.usage,
                         format_string=format_string,
+                        id=message_id,
                     ),
                     inner_messages=inner_messages,
                 )
@@ -1030,9 +1038,13 @@ async def _process_model_result(
                         content=model_result.content,
                         source=agent_name,
                         models_usage=model_result.usage,
+                        id=message_id,
                     ),
                     inner_messages=inner_messages,
                 )
+            # Clean up the correlation map entry to avoid memory leaks
+            if id(model_result) in cls._message_id_map:
+                del cls._message_id_map[id(model_result)]
             return
 
         # Otherwise, we have function calls
@@ -1202,6 +1214,9 @@ async def _reflect_on_tool_use_flow(
 
         reflection_result: Optional[CreateResult] = None
 
+        # Generate a message ID for correlation between chunks and final message
+        message_id = str(uuid.uuid4())
+
         if model_client_stream:
             async for chunk in model_client.create_stream(
                 llm_messages,
@@ -1210,7 +1225,7 @@ async def _reflect_on_tool_use_flow(
                 if isinstance(chunk, CreateResult):
                     reflection_result = chunk
                 elif isinstance(chunk, str):
-                    yield ModelClientStreamingChunkEvent(content=chunk, source=agent_name)
+                    yield ModelClientStreamingChunkEvent(content=chunk, source=agent_name, full_message_id=message_id)
                 else:
                     raise RuntimeError(f"Invalid chunk type: {type(chunk)}")
         else:
@@ -1241,6 +1256,7 @@ async def _reflect_on_tool_use_flow(
                     content=content,
                     source=agent_name,
                     models_usage=reflection_result.usage,
+                    id=message_id,
                 ),
                 inner_messages=inner_messages,
             )
@@ -1250,6 +1266,7 @@ async def _reflect_on_tool_use_flow(
                     content=reflection_result.content,
                     source=agent_name,
                     models_usage=reflection_result.usage,
+                    id=message_id,
                 ),
                 inner_messages=inner_messages,
             )
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/messages.py b/python/packages/autogen-agentchat/src/autogen_agentchat/messages.py
index 715e59ad4f8f..c67062c753d2 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/messages.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/messages.py
@@ -4,6 +4,7 @@
 class and includes specific fields relevant to the type of message being sent.
 """
 
+import uuid
 from abc import ABC, abstractmethod
 from datetime import datetime, timezone
 from typing import Any, Dict, Generic, List, Literal, Mapping, Optional, Type, TypeVar
@@ -77,6 +78,9 @@ class BaseChatMessage(BaseMessage, ABC):
     message using models and return a response as another :class:`BaseChatMessage`.
     """
 
+    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
+    """Unique identifier for this message."""
+
     source: str
     """The name of the agent that sent this message."""
 
@@ -149,6 +153,9 @@ class BaseAgentEvent(BaseMessage, ABC):
     a custom rendering of the content.
     """
 
+    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
+    """Unique identifier for this event."""
+
     source: str
     """The name of the agent that sent this message."""
 
@@ -523,6 +530,10 @@ class ModelClientStreamingChunkEvent(BaseAgentEvent):
     content: str
     """A string chunk from the model client."""
 
+    full_message_id: str | None = None
+    """Optional reference to the complete message that may come after the chunks.
+    This allows consumers of the stream to correlate chunks with the eventual completed message."""
+
     type: Literal["ModelClientStreamingChunkEvent"] = "ModelClientStreamingChunkEvent"
 
     def to_text(self) -> str:
diff --git a/python/packages/autogen-agentchat/tests/test_assistant_agent.py b/python/packages/autogen-agentchat/tests/test_assistant_agent.py
index 219443a21609..5ffb731c1abf 100644
--- a/python/packages/autogen-agentchat/tests/test_assistant_agent.py
+++ b/python/packages/autogen-agentchat/tests/test_assistant_agent.py
@@ -1419,8 +1419,8 @@ class AgentResponse(BaseModel):
     result = await restored_agent.run()
 
     assert isinstance(result.messages[0], StructuredMessage)
-    assert result.messages[0].content.result == "All good"  # type: ignore[reportUnknownMemberType]
-    assert result.messages[0].content.status == "ok"  # type: ignore[reportUnknownMemberType]
+    assert result.messages[0].content.result == "All good"  # type: ignore
+    assert result.messages[0].content.status == "ok"  # type: ignore
 
 
 @pytest.mark.asyncio
diff --git a/python/packages/autogen-agentchat/tests/test_messages.py b/python/packages/autogen-agentchat/tests/test_messages.py
index 17f53f27e782..0e5ff78dd35e 100644
--- a/python/packages/autogen-agentchat/tests/test_messages.py
+++ b/python/packages/autogen-agentchat/tests/test_messages.py
@@ -1,4 +1,5 @@
 import json
+import uuid
 from typing import List
 
 import pytest
@@ -54,7 +55,7 @@ def test_structured_message() -> None:
 
 
 def test_structured_message_component() -> None:
-    # Create a structured message with the test contentformat_string="this is a string {field1} and this is an int {field2}"
+    # Create a structured message with the test content
     format_string = "this is a string {field1} and this is an int {field2}"
     s_m = StructuredMessageFactory(input_model=TestContent, format_string=format_string)
     config = s_m.dump_component()
@@ -187,3 +188,114 @@ def test_union_types() -> None:
     loaded_container = TestContainer.model_validate(data)
     assert loaded_container.chat_messages == chat_messages
     assert loaded_container.agent_events == agent_events
+
+
+def test_message_id_field() -> None:
+    """Test that messages have unique ID fields automatically generated."""
+    # Test BaseChatMessage subclass (TextMessage)
+    message1 = TextMessage(source="test_agent", content="Hello, world!")
+    message2 = TextMessage(source="test_agent", content="Hello, world!")
+
+    # Check that IDs are present and unique
+    assert hasattr(message1, "id")
+    assert hasattr(message2, "id")
+    assert message1.id != message2.id
+    assert isinstance(message1.id, str)
+    assert isinstance(message2.id, str)
+
+    # Check that IDs are valid UUIDs
+    try:
+        uuid.UUID(message1.id)
+        uuid.UUID(message2.id)
+    except ValueError:
+        pytest.fail("Generated IDs are not valid UUIDs")
+
+    # Test BaseAgentEvent subclass (ModelClientStreamingChunkEvent)
+    event1 = ModelClientStreamingChunkEvent(source="test_agent", content="chunk1")
+    event2 = ModelClientStreamingChunkEvent(source="test_agent", content="chunk2")
+
+    # Check that IDs are present and unique
+    assert hasattr(event1, "id")
+    assert hasattr(event2, "id")
+    assert event1.id != event2.id
+    assert isinstance(event1.id, str)
+    assert isinstance(event2.id, str)
+
+    # Check that IDs are valid UUIDs
+    try:
+        uuid.UUID(event1.id)
+        uuid.UUID(event2.id)
+    except ValueError:
+        pytest.fail("Generated IDs are not valid UUIDs")
+
+
+def test_custom_message_id() -> None:
+    """Test that custom IDs can be provided."""
+    custom_id = "custom-message-id-123"
+    message = TextMessage(id=custom_id, source="test_agent", content="Hello, world!")
+
+    assert message.id == custom_id
+
+    custom_event_id = "custom-event-id-456"
+    event = ModelClientStreamingChunkEvent(id=custom_event_id, source="test_agent", content="chunk")
+
+    assert event.id == custom_event_id
+
+
+def test_streaming_chunk_full_message_id() -> None:
+    """Test the full_message_id field in ModelClientStreamingChunkEvent."""
+    # Test without full_message_id
+    chunk1 = ModelClientStreamingChunkEvent(source="test_agent", content="chunk1")
+    assert chunk1.full_message_id is None
+
+    # Test with full_message_id
+    full_msg_id = "full-message-123"
+    chunk2 = ModelClientStreamingChunkEvent(source="test_agent", content="chunk2", full_message_id=full_msg_id)
+    assert chunk2.full_message_id == full_msg_id
+
+    # Test that chunk has its own ID separate from full_message_id
+    assert chunk2.id != chunk2.full_message_id
+    assert isinstance(chunk2.id, str)
+
+    # Verify chunk ID is a valid UUID
+    try:
+        uuid.UUID(chunk2.id)
+    except ValueError:
+        pytest.fail("Chunk ID is not a valid UUID")
+
+
+def test_message_serialization_with_id() -> None:
+    """Test that messages with IDs serialize and deserialize correctly."""
+    # Create a message with auto-generated ID
+    original_message = TextMessage(source="test_agent", content="Hello, world!")
+    original_id = original_message.id
+
+    # Serialize to dict
+    message_data = original_message.model_dump()
+    assert "id" in message_data
+    assert message_data["id"] == original_id
+
+    # Deserialize from dict
+    restored_message = TextMessage.model_validate(message_data)
+    assert restored_message.id == original_id
+    assert restored_message.source == "test_agent"
+    assert restored_message.content == "Hello, world!"
+
+    # Test with streaming chunk event
+    original_chunk = ModelClientStreamingChunkEvent(
+        source="test_agent", content="chunk", full_message_id="full-msg-123"
+    )
+    original_chunk_id = original_chunk.id
+
+    # Serialize to dict
+    chunk_data = original_chunk.model_dump()
+    assert "id" in chunk_data
+    assert "full_message_id" in chunk_data
+    assert chunk_data["id"] == original_chunk_id
+    assert chunk_data["full_message_id"] == "full-msg-123"
+
+    # Deserialize from dict
+    restored_chunk = ModelClientStreamingChunkEvent.model_validate(chunk_data)
+    assert restored_chunk.id == original_chunk_id
+    assert restored_chunk.full_message_id == "full-msg-123"
+    assert restored_chunk.content == "chunk"
diff --git a/python/packages/autogen-agentchat/tests/test_streaming_message_id_correlation.py b/python/packages/autogen-agentchat/tests/test_streaming_message_id_correlation.py
new file mode 100644
index 000000000000..9031c532a6ea
--- /dev/null
+++ b/python/packages/autogen-agentchat/tests/test_streaming_message_id_correlation.py
@@ -0,0 +1,99 @@
+from typing import List, Optional
+
+import pytest
+from autogen_agentchat.agents import AssistantAgent
+from autogen_agentchat.base import TaskResult
+from autogen_agentchat.messages import ModelClientStreamingChunkEvent, TextMessage
+from autogen_core import FunctionCall
+from autogen_core.models import CreateResult, ModelFamily, RequestUsage
+from autogen_ext.models.replay import ReplayChatCompletionClient
+
+
+async def _echo_function(input: str) -> str:
+    return input
+
+
+@pytest.mark.asyncio
+async def test_streaming_message_id_correlation() -> None:
+    """Test that streaming chunks have full_message_id that matches final message ID."""
+    mock_client = ReplayChatCompletionClient(
+        [
+            "Response to message",
+        ]
+    )
+    agent = AssistantAgent(
+        "test_agent",
+        model_client=mock_client,
+        model_client_stream=True,
+    )
+
+    # Track all chunks and the final message
+    chunks: List[ModelClientStreamingChunkEvent] = []
+    final_message: Optional[TextMessage] = None
+
+    async for message in agent.run_stream(task="task"):
+        if isinstance(message, TaskResult):
+            assert len(message.messages) == 2
+            assert isinstance(message.messages[0], TextMessage)
+            assert isinstance(message.messages[1], TextMessage)
+            final_message = message.messages[1]
+        elif isinstance(message, ModelClientStreamingChunkEvent):
+            chunks.append(message)
+
+    # Verify we got chunks and a final message
+    assert len(chunks) > 0
+    assert final_message is not None
+
+    # Every chunk should have the same full_message_id as the final message's id
+    for chunk in chunks:
+        assert chunk.full_message_id == final_message.id
+
+    # Test the reflect_on_tool_use streaming case
+    mock_client = ReplayChatCompletionClient(
+        [
+            CreateResult(
+                content=[
+                    FunctionCall(id="1", name="_echo_function", arguments=r'{"input": "task"}'),
+                ],
+                finish_reason="function_calls",
+                usage=RequestUsage(prompt_tokens=10, completion_tokens=5),
+                cached=False,
+            ),
+            "Example reflection response",
+        ],
+        model_info={
+            "function_calling": True,
+            "vision": False,
+            "json_output": False,
+            "family": ModelFamily.GPT_4,
+            "structured_output": False,
+        },
+    )
+
+    agent = AssistantAgent(
+        "test_agent",
+        model_client=mock_client,
+        model_client_stream=True,
+        reflect_on_tool_use=True,
+        tools=[_echo_function],
+    )
+
+    # Track reflection chunks and final message
+    reflection_chunks: List[ModelClientStreamingChunkEvent] = []
+    final_reflection_message: Optional[TextMessage] = None
+
+    async for message in agent.run_stream(task="task"):
+        if isinstance(message, TaskResult):
+            # The last message should be the reflection result
+            if isinstance(message.messages[-1], TextMessage):
+                final_reflection_message = message.messages[-1]
+        elif isinstance(message, ModelClientStreamingChunkEvent):
+            reflection_chunks.append(message)
+
+    # Verify we got reflection chunks and a final message
+    assert len(reflection_chunks) > 0
+    assert final_reflection_message is not None
+
+    # Every reflection chunk should have the same full_message_id as the final message's id
+    for chunk in reflection_chunks:
+        assert chunk.full_message_id == final_reflection_message.id  # type: ignore


Code review comments:
diff_hunk: @@ -655,6 +646,9 @@ async def run_reasoning_agent() -> None:
     component_config_schema = AssistantAgentConfig
     component_provider_override = "autogen_agentchat.agents.AssistantAgent"
 
+    # Store correlation between CreateResult objects and message IDs
+    _message_id_map: Dict[int, str] = {}
path: python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py
text: @user1:
Pass id as an argument to _call_LLM instead. 

@author:
Hello Sir,
I have implemented the changes as per your suggestions.