Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Draft] Add Example ChromadbMemory in Extensions #5308

Open
wants to merge 2 commits into
base: main
Choose a base branch
from

Conversation

victordibia
Copy link
Collaborator

@victordibia victordibia commented Jan 31, 2025

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

Copy link

codecov bot commented Jan 31, 2025

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 3 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.

@victordibia victordibia changed the title Add Example ChromadbMemory in Extensions [Draft] Add Example ChromadbMemory in Extensions Jan 31, 2025
@EItanya EItanya mentioned this pull request Feb 5, 2025
3 tasks
@gagb
Copy link
Collaborator

gagb commented Feb 13, 2025

@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
Copy link
Collaborator

gagb commented Feb 13, 2025

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
Copy link
Collaborator Author

@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
Copy link
Collaborator

gagb commented Feb 15, 2025

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants