[Draft] Add Example ChromadbMemory in Extensions

[victordibia]

Why are these changes needed?

Shows an example of how to use the Memory interface to implement a just-in-time vector memory based on chromadb.

from autogen_core.memory import MemoryContent, MemoryMimeType
from autogen_agentchat.agents import AssistantAgent
from autogen_ext.models.openai import OpenAIChatCompletionClient
from autogen_agentchat.ui import Console
from autogen_ext.memory.chromadb import ChromaMemory, ChromaMemoryConfig, ChromaMemoryContent

async def main():
    # Initialize ChromaDB memory with custom config
    user_memory = ChromaMemory(
        name="user_preferences",
        config=ChromaMemoryConfig(
            collection_name="preferences",
            persistence_path="./chroma_db",  # Optional: persist to disk
            k=2,  # Return top 5 results
            score_threshold=0.4,  # Minimum similarity score
        )
    )

    # Add user preferences to memory
    await user_memory.add(
        MemoryContent(
            content="The weather should be in metric units",
            mime_type=MemoryMimeType.TEXT,
            metadata={"category": "preferences", "type": "units"}
        )
    )
    
    await user_memory.add(
        MemoryContent(
            content="Meal recipe must be vegan",
            mime_type=MemoryMimeType.TEXT,
            metadata={"category": "preferences", "type": "dietary"}
        )
    )

    async def get_weather(city: str, units: str = "imperial") -> str:
        if units == "imperial":
            return f"The weather in {city} is 73 °F and Sunny."
        elif units == "metric":
            return f"The weather in {city} is 23 °C and Sunny."
        else:
            return f"Sorry, I don't know the weather in {city}."

    # Create assistant agent with ChromaDB memory
    assistant_agent = AssistantAgent(
        name="assistant_agent",
        model_client=OpenAIChatCompletionClient(
            model="gpt-4o",
        ),
        tools=[get_weather],
        memory=[user_memory],
    )

    stream = assistant_agent.run_stream(task="What is the weather in New York?")
    await Console(stream)
 
    await user_memory.close()

# if __name__ == "__main__":
#     import asyncio
#     asyncio.run(main())

Related issue number

Checks

• 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.
• I've added tests (if relevant) corresponding to the changes introduced in this PR.
• I've made sure all auto checks have passed.

[codecov]

Codecov Report

Attention: Patch coverage is 75.73529% with 33 lines in your changes missing coverage. Please review.

Project coverage is 70.65%. Comparing base (a673957) to head (22dd6b0).
Report is 135 commits behind head on main.

Files with missing lines	Patch %	Lines
...ges/autogen-ext/src/autogen_ext/memory/chromadb.py	75.73%	33 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #5308      +/-   ##
==========================================
+ Coverage   70.59%   70.65%   +0.05%     
==========================================
  Files         180      181       +1     
  Lines       11668    11804     +136     
==========================================
+ Hits         8237     8340     +103     
- Misses       3431     3464      +33     

Flag	Coverage Δ
unittests	70.65% <75.73%> (+0.05%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[gagb]

@rickyloynd-microsoft and I discussed this PR and PR #5227.

As a user I need:
clarity -- which memory interface to use when
consistency -- no concept/code duplication

We tried the ListMemory() interface in core and the metadata field in it is useful for storing (task, insight) tuples. That replicated some of the features in #5227 but not all. We should improve clarity about this and consistency between the two APIs.

cc: @ekzhu @jackgerrits

[gagb]

Here is the script we tried. cc @rickyloynd-microsoft

from autogen_agentchat.agents import AssistantAgent
from autogen_agentchat.ui import Console
from autogen_core.memory import ListMemory, MemoryContent, MemoryMimeType, MemoryQueryResult, UpdateContextResult
from autogen_ext.models.openai import OpenAIChatCompletionClient
from autogen_core.model_context import ChatCompletionContext
from autogen_core.models import SystemMessage


class CustomListMemory(ListMemory):


     async def update_context(
        self,
        model_context: ChatCompletionContext,
    ) -> UpdateContextResult:
        """Update the model context by appending memory content.

        This method mutates the provided model_context by adding all memories as a
        SystemMessage.

        Args:
            model_context: The context to update. Will be mutated if memories exist.

        Returns:
            UpdateContextResult containing the memories that were added to the context
        """

        if not self._contents:
            return UpdateContextResult(memories=MemoryQueryResult(results=[]))

        memory_strings = [f"{i}. {str(memory.content) + str(memory.metadata['plan']) if memory.metadata else ''}" for i, memory in enumerate(self._contents, 1)]

        if memory_strings:
            memory_context = "\nRelevant memory content (in chronological order):\n" + "\n".join(memory_strings) + "\n"
            await model_context.add_message(SystemMessage(content=memory_context))

        return UpdateContextResult(memories=MemoryQueryResult(results=self._contents))



async def main():
    # Initialize user memory
    # user_memory = ListMemory()

    user_memory = CustomListMemory()

    # Add user preferences to memory
    await user_memory.add(MemoryContent(content="plan for ordering peripherals from favorite website",
    metadata={"plan": "go to amazon.com > search for peripherals > order"}, 
    mime_type=MemoryMimeType.TEXT))

    await user_memory.add(MemoryContent(content="order meal from grubhub",
    metadata={"plan": "go to grubhub.com > order meal"},
    mime_type=MemoryMimeType.TEXT))

    assistant_agent = AssistantAgent(
        name="assistant_agent",
        model_client=OpenAIChatCompletionClient(
            model="gpt-4o-2024-08-06",
        ),
        memory=[user_memory],
    )

    # Run the agent with a task.
    stream = assistant_agent.run_stream(task="Create a plan to order a keyboard from my favorite website?")
    await Console(stream)

if __name__ == "__main__":
    import asyncio

    asyncio.run(main())

[victordibia]

@gagb , thanks for the excellent example above - it indeed is in the spirit of the Memory interface. It would be great to figure out what is needed to cover functionality that is left.

Is your question on consistency related when do we use the Memory interface vs Task Centric Memory #5227 ?
I think a good way to look at it is the Memory is a base interface we provide so that anyone can inherit from it and built Memory implementations that work with the framework (and example of this being compatibility with AgentChat AssistantAgent like your example shows).

If you'd like your Memory implementation pluggable across the ecosystem, inherit from Memory

What do you think?

[gagb]

Yeah I agree with you.

I think it would be great if Ricky's PR was using this interface.

And if that is not possible there need to be a clear reason.

[husseinmozannar]

@victordibia thanks for the PR! I noticed it is still in draft, are you planning on more features or are you waiting on further review/testing? If so, I can help with the later.

[victordibia]

I recall having some mypy/pyright errors I needed to address. Will get back to it this today/tomorrow and share where I need help, thanks!