Task-Centric Memory

[rickyloynd-microsoft]

(EXPERIMENTAL RESEARCH IN PROGRESS)

In 2023 AutoGen introduced Teachable Agents that users could teach new facts, preferences and skills. But teachable agents were limited in several ways: They could only be ConversableAgent subclasses, they couldn't learn a new skill unless the user stated (in a single turn) both the task and how to solve it, and they couldn't learn on their own. Task-Centric Memory overcomes these limitations, allowing users to teach arbitrary agents (or teams) more flexibly and reliably, and enabling agents to learn from their own trial-and-error experiences.

This PR is large and complex. All of the files are new, and most of the added components depend on the others to run at all. But the review process can be accelerated if approached in the following order.

1. Start with the Task-Centric Memory README.
   1. Install the memory extension locally, since it won't be in pypi until it's merged. In the agentic_memory branch, and the python/packages directory:
      • pip install -e autogen-agentchat
      • pip install -e autogen-ext[openai]
      • pip install -e autogen-ext[task-centric-memory]
   2. Run the Quickstart sample code, then immediately open the ~/pagelogs/quick/0 Call Tree.html file in a browser to view the work in progress.
   3. Click through the web page links to see the details.
2. Continue through the rest of the main README to get a high-level overview of the architecture.
3. Read through the code samples README, running each of the 4 code samples while viewing their page logs.
4. Skim through the 4 code samples, along with their corresponding yaml config files:
   1. eval_retrieval.py
   2. eval_teachability.py
   3. eval_learning_from_demonstration.py
   4. eval_self_teaching.py
5. Read task_centric_memory_controller.py, referring back to the previously generated page logs as needed. This is the most important and complex file in the PR.
6. Read the remaining core files.
   1. _task_centric_memory_bank.py
   2. _string_similarity_map.py
   3. _prompter.py
7. Read the supporting files in the utils dir.
   1. apprentice.py
   2. grader.py
   3. page_logger.py
   4. _functions.py

[rickyloynd-microsoft]

Thanks for the PR, Ricky. I left a few comments.

This PR is very big. And it's hard to review these many number of file changes. Is it possible to make this PR incrementally?

I think that would help everyone to give you feedback more carefully and on the most important bits first. But I also realize that it may not be feasible to break this PR.

Followed your suggestion to put a recommended list of steps in the PR intro to help reviewers.

[ekzhu]

Another example perhaps like this. It can be just another code block in this file for now. In a future PR when we write a user guide for this memory module, we can move this example to the user guide.

import asyncio
from dataclasses import dataclass
from typing import List

from autogen_core import AgentId, MessageContext, RoutedAgent, SingleThreadedAgentRuntime, message_handler
from autogen_core.models import ChatCompletionClient, LLMMessage, SystemMessage, UserMessage
from autogen_ext.models.openai import OpenAIChatCompletionClient
from autogen_ext.task_centric_memory import PageLogger, TaskCentricMemoryController


@dataclass
class Message:
    content: str


class MemoryEnabledAgent(RoutedAgent):
    def __init__(
        self, description: str, model_client: ChatCompletionClient, task_memory_controller: TaskCentricMemoryController
    ) -> None:
        super().__init__(description)
        self._model_client = model_client
        self._task_memory_controller = task_memory_controller

    @message_handler
    async def handle_message(self, message: Message, context: MessageContext) -> Message:
        # Retrieve relevant memories for the task.
        memos = await self._task_memory_controller.retrieve_relevant_memos(task=message.content)
        # Format the memories for the model.
        formatted_memos = "Relevant context about the user: \n\n" + "\n".join([memo.insight for memo in memos])
        print(f"{'-'*38}Memos{'-'*38}:\n{formatted_memos}\n{'-'*80}")
        # Create the messages for the model with the retrieved memories.
        messages: List[LLMMessage] = [
            SystemMessage(content="You are a helpful assistant."),
            UserMessage(content=formatted_memos, source="user"),
            UserMessage(content=message.content, source="user"),
        ]
        # Call the model with the messages.
        model_result = await self._model_client.create(messages=messages)
        assert isinstance(model_result.content, str)
        # Send the model's response to the user.
        return Message(content=model_result.content)


async def main() -> None:
    client = OpenAIChatCompletionClient(model="gpt-4o")
    page_logger = PageLogger(config={"level": "DEBUG", "path": "~/pagelogs/quickstart"})  # Optional, but very useful.
    memory_controller = TaskCentricMemoryController(reset=True, client=client, logger=page_logger)

    # Add a few task-insight pairs as memories, where an insight can be any string that may help solve the task.
    await memory_controller.add_memo(task="What color do I like?", insight="Deep blue is my favorite color")
    await memory_controller.add_memo(task="What's another color I like?", insight="I really like cyan")
    await memory_controller.add_memo(task="What's my favorite food?", insight="Halibut is my favorite")

    # Create an agent runtime.
    runtime = SingleThreadedAgentRuntime()

    # Start the agent runtime.
    runtime.start()

    # Register the agent type.
    await MemoryEnabledAgent.register(
        runtime,
        "memory_enabled_agent",
        lambda: MemoryEnabledAgent(
            "A agent with memory", model_client=client, task_memory_controller=memory_controller
        ),
    )

    # Send a direct message to the agent.
    response = await runtime.send_message(
        Message(content="What colors do I like most"), AgentId("memory_enabled_agent", "default")
    )

    print("Agent response: " + response.content)

    # Stop the agent runtime.
    await runtime.stop()


asyncio.run(main())

--------------------------------------Memos--------------------------------------:
Relevant context about the user: 

Deep blue is my favorite color
I really like cyan
--------------------------------------------------------------------------------
Agent response: Based on the information you've provided, you really like deep blue and cyan.

To show how to use this as part of an agent implementation in the Core API.

[ekzhu]

@victordibia could you comment how can we perhaps leverage this module as an implementation of the autogen_core.memory.Memory interface?

Not for this PR but for future one.

[victordibia]

Good idea.
It seems to me that the example above could be indeed implemented (to some extent) using the Memory interface.

• add maps to memory_controller.add_memo
• and update_context maps to formatted_memos = ...

[ekzhu]

Yeah. I do think using the controller we can implement the autogen_core.memory.Memory interface. I can give it a try.

If that's the case we should move this module to autogen_ext.memory.task_centric_memory.

[rickyloynd-microsoft]

Without losing any of the current functionality of TaskCentricMemoryController? Like learning from its own experience?? (That's TaskCentricMemoryController.train_on_task(self, task: str, expected_answer: str), which requires the expected answer.)

TaskCentricMemoryController is still active research so the interface is highly likely to change, as we discussed before. Why do you now want it in Core?

[ekzhu]

Why do you now want it in Core?

We are not moving it to core. It's staying in extension. My message earlier refers to autogen_ext module.

I am just trying to play with it and see what it takes me.

[rickyloynd-microsoft]

New code example from @ekzhu added to the README.

[rickyloynd-microsoft]

Yeah. I do think using the controller we can implement the autogen_core.memory.Memory interface. I can give it a try.

If that's the case we should move this module to autogen_ext.memory.task_centric_memory.

@ekzhu, should we wait for the results of your attempt to use TaskCentricMemoryController to implement the autogen_core.memory.Memory interface before we try to clarify the relations and connections between TaskCentricMemory and the autogen_core.memory.Memory interface?

[ekzhu]

Add an example code block here for the simplest example that you have in the README.md and then add links to the sample directory for more advanced samples.

See https://github.com/microsoft/autogen/blob/agentic_memory/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py#L1258 for an example of how to add doc-string code block.

[rickyloynd-microsoft]

Added

[husseinmozannar]

as discussed, there is an option to combine these two lines in a single API call, it would sacrifice accuracy potentially, but go from minimum 2 LLM calls to 1 LLM call

[husseinmozannar]

there is potential to combine these three LLM calls into a single call, it depends on the LLM, but GPT-4o should be able to do this in a single shot with a COT prompt, other less capable models might need extra reflection