From dd0781a76bb98e594fc6df23586661973e327cde Mon Sep 17 00:00:00 2001
From: Victor Dibia <victordibia@microsoft.com>
Date: Wed, 26 Feb 2025 10:40:00 -0800
Subject: [PATCH] Add Serialization Instruction for  MemoryContent (#5727)

<!-- 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?

- Running a team generates a MemoryQueryEvent has a MemoryContent field
and is part of the model context
- Saving team state (`team.save_state()`) includes serializing model
context
- MemoryContent has a mime_type field, which was not being properly
serialized.
"Object of type MemoryMimeType is not JSON serializable"

This PR? -> add explicit serialization instruction for the mimetype
field.

```python
import asyncio
import logging
import json
import yaml, aiofiles
from autogen_agentchat.agents import AssistantAgent
from autogen_agentchat.teams import RoundRobinGroupChat
from autogen_ext.models.openai import OpenAIChatCompletionClient
from autogen_agentchat.conditions import MaxMessageTermination
from autogen_core.memory import ListMemory, MemoryContent, MemoryMimeType
from autogen_agentchat.ui import Console

logger = logging.getLogger(__name__)

state_path = "team_state.json"

model_client =  OpenAIChatCompletionClient(model="gpt-4o-mini")

max_msg_termination = MaxMessageTermination(max_messages=3)
# Initialize user memory
user_memory = ListMemory()

# Add user preferences to memory
await user_memory.add(MemoryContent(content="The weather should be in metric units", mime_type=MemoryMimeType.TEXT))

# Create the team.
agent = AssistantAgent(
    name="assistant",
    model_client=model_client,
    system_message="You are a helpful assistant.",
    memory = [user_memory],
)
yoda = AssistantAgent(
    name="yoda",
    model_client=model_client,
    system_message="Repeat the same message in the tone of Yoda.",
)

team = RoundRobinGroupChat(
    [agent, yoda],
    termination_condition=max_msg_termination
)

await Console(team.run_stream(task="Hi, How are you ?"))

# Save team state to file.
state = await team.save_state()
with open(state_path, "w") as f:
    json.dump(state, f)
# Load team state from file.
with open("team_state.json", "r") as f:
    team_state = json.load(f)

await team.load_state(state)
await Console( team.run_stream(task="What was the last thing that was said in this conversation "))
```

<!-- Please give a short summary of the change and the problem this
solves. -->

## Related issue number

<!-- For example: "Closes #1234" -->

Closes #5688
## Checks

- [ ] 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.
- [ ] I've added tests (if relevant) corresponding to the changes
introduced in this PR.
- [ ] I've made sure all auto checks have passed.
---
 .../autogen-core/src/autogen_core/memory/_base_memory.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/python/packages/autogen-core/src/autogen_core/memory/_base_memory.py b/python/packages/autogen-core/src/autogen_core/memory/_base_memory.py
index 577c23fda..5385b37d3 100644
--- a/python/packages/autogen-core/src/autogen_core/memory/_base_memory.py
+++ b/python/packages/autogen-core/src/autogen_core/memory/_base_memory.py
@@ -2,7 +2,7 @@ from abc import ABC, abstractmethod
 from enum import Enum
 from typing import Any, Dict, List, Union
 
-from pydantic import BaseModel, ConfigDict
+from pydantic import BaseModel, ConfigDict, field_serializer
 
 from .._cancellation_token import CancellationToken
 from .._component_config import ComponentBase
@@ -37,6 +37,13 @@ class MemoryContent(BaseModel):
 
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
+    @field_serializer("mime_type")
+    def serialize_mime_type(self, mime_type: MemoryMimeType | str) -> str:
+        """Serialize the MIME type to a string."""
+        if isinstance(mime_type, MemoryMimeType):
+            return mime_type.value
+        return mime_type
+
 
 class MemoryQueryResult(BaseModel):
     """Result of a memory :meth:`~autogen_core.memory.Memory.query` operation."""