State Management does not work with OpenAIAssistantAgent

[ssbnord]

What happened?

This tutorial works for AssistantAgent not for OpenAIAssistantAgent. It looks like the save_state and load_state methods use the default impl here which return

{'type': 'BaseState', 'version': '1.0.0'}

It works as expected for AssistantAgent.

What did you expect to happen?

State management should work for OpenAIAssistantAgent similar to AssistantAgent.

How can we reproduce it (as minimally and precisely as possible)?

import asyncio
from autogen_ext.agents.openai import OpenAIAssistantAgent
from autogen_core import CancellationToken

from autogen_agentchat.messages import TextMessage
from openai import AsyncAzureOpenAI
import os


async def example():

    async_model_client = AsyncAzureOpenAI(
        api_key = os.environ.get("API_KEY"),
        azure_endpoint = os.environ.get("API_BASE"),
        api_version= "2024-07-01-preview",
    )

    agent = OpenAIAssistantAgent(name="FunnyAgent",
                            description="Tells jokes",
                            client=async_model_client,
                            model="gpt-4o",
                            instructions="You tell jokes.",
                        )

    response = await agent.on_messages(
        [TextMessage(source="user", content="Tell me a joke")], CancellationToken()
    )

    print(response.chat_message.content)

    agent_state = await agent.save_state()
    print(agent_state)

asyncio.run(example())

Observe that the agent state is returned as

{'type': 'BaseState', 'version': '1.0.0'}

AutoGen version

0.4.2

Which package was this bug in

Extensions

Model used

gpt-4o

Python version

3.12.8

Operating system

MacOS 15.2

Any additional info you think would be helpful for fixing this bug

No response

[ekzhu]

Would love some help on this.

It should be simple as the state should be just the thread id and assistant id. Are there other references we need?

[lspinheiro]

I can pick it up. I will have a look tomorrow.

[ssbnord]

Is it possible to use save_state and load_state so that we can have feature parity between AssistantAgent and GPTAssistantAgent in 0.2. While we can use assistant id and thread id, it requires maintaining that state in another store.

[lspinheiro]

I have a PR here: #5352

@ssbnord , please have a look and see if that works for you.