Task-Centric Memory

python/.gitignore

@@ -157,7 +157,7 @@ cython_debug/
 #  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
-#.idea/
+.idea/
 
 .ruff_cache/
 

python/packages/autogen-ext/pyproject.toml

@@ -66,6 +66,8 @@ jupyter-executor = [
     "nbclient>=0.10.2",
 ]
 
+agentic-memory = ["chromadb>=0.6.3"]
+
 semantic-kernel-core = [
     "semantic-kernel>=1.17.1",
 ]
@@ -140,7 +142,7 @@ dev = [
 [tool.ruff]
 extend = "../../pyproject.toml"
 include = ["src/**", "tests/*.py"]
-exclude = ["src/autogen_ext/agents/web_surfer/*.js", "src/autogen_ext/runtimes/grpc/protos", "tests/protos"]
+exclude = ["src/autogen_ext/agents/web_surfer/*.js", "src/autogen_ext/runtimes/grpc/protos", "tests/protos", "README.md"]
 
 [tool.pyright]
 extends = "../../pyproject.toml"

python/packages/autogen-ext/src/autogen_ext/agentic_memory/README.md

@@ -0,0 +1,81 @@
+# Agentic Memory
+
+This AutoGen extension provides an implementation of agentic memory, which we define as a
+broad ability for AI agents to accomplish tasks more effectively by learning quickly and continually (over the long term).
+This is distinct from what RAG or long context windows can provide.
+While still under active research and development, this implementation of agentic memory
+can be attached to virtually any unmodified AI agent, and is designed to enable agents that:
+
+* Remember guidance, corrections, and demonstrations provided by users.
+* Succeed more frequently on tasks after finding successful solutions to similar tasks.
+* Learn and adapt quickly to changing circumstances to enable workflows that are dynamic and self-healing.
+
+The implementation is also intended to:
+
+* Be general purpose, unconstrained by types and schemas required by standard databases.
+* Augment rather than interfere with an agent’s special capabilities, such as powerful reasoning, long-horizon autonomy, and tool handling.
+* Operate in both foreground and background modes, so that an agent can discuss tasks with a user (in the foreground)
+then work productively on those tasks (in the background) while the user does other things.
+* Allow for fine-grained transparency and auditing of individual memories by human users or other agents.
+* Allow agents to be personalized (to a single user) as well as specialized (to a subject, domain or project).
+The benefits of personalization scale linearly with the number of users, but the benefits of domain specialization
+can scale quadratically with the number of users working in that domain, as insights gained from interactions with one user
+can benefit other users in similar situations.
+* Support multiple memory banks dynamically attached to an agent at runtime.
+* Enable enforcement of security boundaries at the level of individual memory banks.
+* Allow users to download and port memory banks between agents and systems.
+
+![agentic_memory.png](../../../imgs/agentic_memory.png)
+
+The block diagram above outlines the key components of our baseline agentic memory architecture,
+which augments an agent or team with agentic memory mechanisms.
+
+The **Agentic Memory Controller** implements the fast-learning methods described below,
+and manages communication with an **Agentic Memory Bank** containing a vector DB and associated structures.
+
+The **Apprentice** is a minimal reference implementation that wraps the combination of agentic memory plus some agent or team.
+Certain applications will use the Apprentice,
+while others will directly instantiate and call the Agentic Memory Controller.
+
+We’ve successfully tested agentic memory with a simple AssistantAgent and MagenticOneGroupChat.
+
+## Memory Creation and Storage
+
+Each stored memory (called a _memo_) is an insight (in text form) crafted to help the agent accomplish future tasks that are similar
+to some task encountered in the past. If the user provides advice for solving a given task,
+the advice is extracted and stored as an insight. If the user demonstrates how to perform a task,
+the task and demonstration are stored together as an insight that could be applied to similar but different tasks.
+If the agent is given a task (free of side-effects) and some means of determining success or failure,
+the memory controller repeats the following learning loop in the background some number of times:
+
+1. Test the agent on the task a few times to check for a failure.
+2. If a failure is found, analyze the agent’s response in order to:
+   1. Diagnose the failure of reasoning or missing information,
+   2. Phrase a general piece of advice, such as what a teacher might give to a student,
+   3. Temporarily append this advice to the task description,
+   4. Return to step 1.
+   5. If some piece of advice succeeds in helping the agent solve the task a number of times, add the advice as an insight to memory.
+3. For each insight to be stored in memory, an LLM is prompted to generate a set of free-form, multi-word topics related to the insight. Each topic is embedded to a fixed-length vector and stored in a vector DB mapping it to the topic’s related insight.
+
+## Memory Retrieval and Usage
+
+When the agent is given a task, the following steps are performed by the memory controller:
+1. The task is rephrased into a generalized form.
+2. A set of free-form, multi-word query topics are generated from the generalized task.
+3. A potentially large number of previously stored topics, those most similar to each query topic, are retrieved from the vector DB along with the insights they map to.
+4. These candidate memos are filtered by the aggregate similarity of their stored topics to the query topics.
+5. In the final filtering stage, an LLM is prompted to return only those insights that seem potentially useful in solving the task at hand.
+
+Retrieved insights that pass the filtering steps are listed under a heading like
+“Important insights that may help solve tasks like this”, then appended to the task description before it is passed to the agent as usual.
+
+## Setup and Usage
+
+Install AutoGen and its extension package as follows:
+
+`pip install "autogen-ext[agentic-memory]"`
+
+We provide [sample code](../../../../../samples/agentic_memory) to illustrate the following forms of memory-based fast learning:
+* Agent learning from user advice and corrections
+* Agent learning from user demonstrations
+* Agent learning from its own experience

python/packages/autogen-ext/src/autogen_ext/agentic_memory/__init__.py

@@ -0,0 +1,6 @@
+from .agentic_memory_controller import AgenticMemoryController
+from .apprentice import Apprentice
+from .grader import Grader
+from .page_logger import PageLogger
+
+__all__ = ["Apprentice", "PageLogger", "Grader", "AgenticMemoryController"]

python/packages/autogen-ext/src/autogen_ext/agentic_memory/_agentic_memory_bank.py

@@ -0,0 +1,198 @@
+import os
+import pickle
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple
+
+from ._string_similarity_map import StringSimilarityMap
+from .page_logger import PageLogger
+
+
+@dataclass
+class Memo:
+    """
+    Represents an atomic unit of memory that can be stored in a memory bank and later retrieved.
+    """
+
+    task: str | None  # The task description, if any.
+    insight: str  # A hint, solution, plan, or any other text that may help solve a similar task.
+
+
+class AgenticMemoryBank:
+    """
+    Stores task-completion insights as memories in a vector DB for later retrieval.
+
+    Args:
+        - reset: True to clear the DB before starting.
+        - config: An optional dict that can be used to override the following values:
+            - path: The path to the directory where the memory bank files are stored.
+            - relevance_conversion_threshold: The threshold used to normalize relevance.
+            - n_results: The maximum number of most relevant results to return for any given topic.
+            - distance_threshold: The maximum string-pair distance for a memo to be retrieved.
+        - logger: An optional logger. If None, no logging will be performed.
+
+    Methods:
+        - reset: Forces immediate deletion of all contents, in memory and on disk.
+        - save_memos: Saves the current memo structures (possibly empty) to disk.
+        - contains_memos: Returns True if the memory bank contains any memos.
+        - add_memo: Adds a memo to the memory bank, given topics related to the insight, and optionally the task.
+        - add_task_with_solution: Adds a task-insight pair to the memory bank, to be retrieved together later.
+        - get_relevant_memos: Returns any memos from the memory bank that appear sufficiently relevant to the input topics.
+    """
+
+    def __init__(
+        self,
+        reset: bool,
+        config: Dict[str, Any] | None = None,
+        logger: PageLogger | None = None,
+    ) -> None:
+        if logger is None:
+            logger = PageLogger()  # Nothing will be logged by this object.
+        self.logger = logger
+        self.logger.enter_function()
+
+        # Assign default values that can be overridden by config.
+        memory_dir_path = os.path.expanduser("~/agentic_memory/temp")
+        self.relevance_conversion_threshold = 1.7
+        self.n_results = 25
+        self.distance_threshold = 100
+
+        if config is not None:
+            # Apply any overrides from the config.
+            for key in config:
+                if key == "path":
+                    memory_dir_path = os.path.expanduser(config[key])
+                elif key == "relevance_conversion_threshold":
+                    self.relevance_conversion_threshold = config[key]
+                elif key == "n_results":
+                    self.n_results = config[key]
+                elif key == "distance_threshold":
+                    self.distance_threshold = config[key]
+                else:
+                    self.logger.error('Unexpected item in config: ["{}"] = {}'.format(key, config[key]))
+
+        path_to_db_dir = os.path.join(memory_dir_path, "string_map")
+        self.path_to_dict = os.path.join(memory_dir_path, "uid_memo_dict.pkl")
+
+        self.string_map = StringSimilarityMap(reset=reset, path_to_db_dir=path_to_db_dir, logger=self.logger)
+
+        # Load or create the associated memo dict on disk.
+        self.uid_memo_dict: Dict[str, Memo] = {}
+        self.last_memo_id = 0
+        if (not reset) and os.path.exists(self.path_to_dict):
+            self.logger.info("\nLOADING MEMOS FROM DISK  {}".format(self.path_to_dict))
+            self.logger.info("    Location = {}".format(self.path_to_dict))
+            with open(self.path_to_dict, "rb") as f:
+                self.uid_memo_dict = pickle.load(f)
+                self.last_memo_id = len(self.uid_memo_dict)
+                self.logger.info("\n{} MEMOS LOADED".format(len(self.uid_memo_dict)))
+
+        # Clear the DB if requested.
+        if reset:
+            self._reset_memos()
+
+        self.logger.leave_function()
+
+    def reset(self) -> None:
+        """
+        Forces immediate deletion of all contents, in memory and on disk.
+        """
+        self.string_map.reset_db()
+        self._reset_memos()
+
+    def _reset_memos(self) -> None:
+        """
+        Forces immediate deletion of the memos, in memory and on disk.
+        """
+        self.uid_memo_dict = {}
+        self.save_memos()
+
+    def save_memos(self) -> None:
+        """
+        Saves the current memo structures (possibly empty) to disk.
+        """
+        self.string_map.save_string_pairs()
+        with open(self.path_to_dict, "wb") as file:
+            pickle.dump(self.uid_memo_dict, file)
+
+    def contains_memos(self) -> bool:
+        """
+        Returns True if the memory bank contains any memo.
+        """
+        return len(self.uid_memo_dict) > 0
+
+    def _map_topics_to_memo(self, topics: List[str], memo_id: str, memo: Memo) -> None:
+        """
+        Adds a mapping in the vec DB from each topic to the memo.
+        """
+        self.logger.enter_function()
+        self.logger.info("\nINSIGHT\n{}".format(memo.insight))
+        for topic in topics:
+            self.logger.info("\n TOPIC = {}".format(topic))
+            self.string_map.add_input_output_pair(topic, memo_id)
+        self.uid_memo_dict[memo_id] = memo
+        self.logger.leave_function()
+
+    def add_memo(self, insight_str: str, topics: List[str], task_str: Optional[str] = None) -> None:
+        """
+        Adds an insight to the memory bank, given topics related to the insight, and optionally the task.
+        """
+        self.logger.enter_function()
+        self.last_memo_id += 1
+        id_str = str(self.last_memo_id)
+        insight = Memo(insight=insight_str, task=task_str)
+        self._map_topics_to_memo(topics, id_str, insight)
+        self.logger.leave_function()
+
+    def add_task_with_solution(self, task: str, solution: str, topics: List[str]) -> None:
+        """
+        Adds a task-solution pair to the memory bank, to be retrieved together later as a combined insight.
+        This is useful when the insight is a demonstration of how to solve a given type of task.
+        """
+        self.logger.enter_function()
+        self.last_memo_id += 1
+        id_str = str(self.last_memo_id)
+        # Prepend the insight to the task description for context.
+        insight_str = "Example task:\n\n{}\n\nExample solution:\n\n{}".format(task, solution)
+        memo = Memo(insight=insight_str, task=task)
+        self._map_topics_to_memo(topics, id_str, memo)
+        self.logger.leave_function()
+
+    def get_relevant_memos(self, topics: List[str]) -> List[Memo]:
+        """
+        Returns any memos from the memory bank that appear sufficiently relevant to the input topics.
+        """
+        self.logger.enter_function()
+
+        # Retrieve all topic matches, and gather them into a single list.
+        matches: List[Tuple[str, str, float]] = []  # Each match is a tuple: (topic, memo_id, distance)
+        for topic in topics:
+            matches.extend(self.string_map.get_related_string_pairs(topic, self.n_results, self.distance_threshold))
+
+        # Build a dict of memo-relevance pairs from the matches.
+        memo_relevance_dict: Dict[str, float] = {}
+        for match in matches:
+            relevance = self.relevance_conversion_threshold - match[2]
+            memo_id = match[1]
+            if memo_id in memo_relevance_dict:
+                memo_relevance_dict[memo_id] += relevance
+            else:
+                memo_relevance_dict[memo_id] = relevance
+
+        # Log the details of all the retrieved memos.
+        self.logger.info("\n{} POTENTIALLY RELEVANT MEMOS".format(len(memo_relevance_dict)))
+        for memo_id, relevance in memo_relevance_dict.items():
+            memo = self.uid_memo_dict[memo_id]
+            details = ""
+            if memo.task is not None:
+                details += "\n  TASK: {}\n".format(memo.task)
+            details += "\n  INSIGHT: {}\n\n  RELEVANCE: {:.3f}\n".format(memo.insight, relevance)
+            self.logger.info(details)
+
+        # Compose the list of sufficiently relevant memos to return.
+        memo_list: List[Memo] = []
+        for memo_id in memo_relevance_dict:
+            if memo_relevance_dict[memo_id] >= 0:
+                memo_list.append(self.uid_memo_dict[memo_id])
+
+        self.logger.leave_function()
+        return memo_list