Adding autogen cookbook to help provide options of integration with a multi-agent framework (#1908)

cookbooks/helper/mem0_teachability.py

@@ -0,0 +1,169 @@
+# Copyright (c) 2023 - 2024, Owners of https://github.com/autogen-ai
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+# Portions derived from  https://github.com/microsoft/autogen are under the MIT License.
+# SPDX-License-Identifier: MIT
+# forked from autogen.agentchat.contrib.capabilities.teachability.Teachability
+
+from typing import Dict, Optional, Union
+from autogen.agentchat.assistant_agent import ConversableAgent
+from autogen.agentchat.contrib.capabilities.agent_capability import AgentCapability
+from autogen.agentchat.contrib.text_analyzer_agent import TextAnalyzerAgent
+from termcolor import colored
+from mem0 import Memory
+from mem0.configs.base import MemoryConfig
+
+class Mem0Teachability(AgentCapability):
+    def __init__(
+        self,
+        verbosity: Optional[int] = 0,
+        reset_db: Optional[bool] = False,
+        recall_threshold: Optional[float] = 1.5,
+        max_num_retrievals: Optional[int] = 10,
+        llm_config: Optional[Union[Dict, bool]] = None,
+        agent_id: Optional[str] = None,
+        memory_client: Optional[Memory] = None,
+    ):
+        self.verbosity = verbosity
+        self.recall_threshold = recall_threshold
+        self.max_num_retrievals = max_num_retrievals
+        self.llm_config = llm_config
+        self.analyzer = None
+        self.teachable_agent = None
+        self.agent_id = agent_id
+        self.memory = memory_client if memory_client else Memory()
+
+        if reset_db:
+            self.memory.reset()
+
+    def add_to_agent(self, agent: ConversableAgent):
+        self.teachable_agent = agent
+        agent.register_hook(hookable_method="process_last_received_message", hook=self.process_last_received_message)
+
+        if self.llm_config is None:
+            self.llm_config = agent.llm_config
+        assert self.llm_config, "Teachability requires a valid llm_config."
+
+        self.analyzer = TextAnalyzerAgent(llm_config=self.llm_config)
+
+        agent.update_system_message(
+            agent.system_message
+            + "\nYou've been given the special ability to remember user teachings from prior conversations."
+        )
+
+    def process_last_received_message(self, text: Union[Dict, str]):
+        expanded_text = text
+        if self.memory.get_all(agent_id=self.agent_id):
+            expanded_text = self._consider_memo_retrieval(text)
+        self._consider_memo_storage(text)
+        return expanded_text
+
+    def _consider_memo_storage(self, comment: Union[Dict, str]):
+        memo_added = False
+        response = self._analyze(
+            comment,
+            "Does any part of the TEXT ask the agent to perform a task or solve a problem? Answer with just one word, yes or no.",
+        )
+
+        if "yes" in response.lower():
+            advice = self._analyze(
+                comment,
+                "Briefly copy any advice from the TEXT that may be useful for a similar but different task in the future. But if no advice is present, just respond with 'none'.",
+            )
+
+            if "none" not in advice.lower():
+                task = self._analyze(
+                    comment,
+                    "Briefly copy just the task from the TEXT, then stop. Don't solve it, and don't include any advice.",
+                )
+
+                general_task = self._analyze(
+                    task,
+                    "Summarize very briefly, in general terms, the type of task described in the TEXT. Leave out details that might not appear in a similar problem.",
+                )
+
+                if self.verbosity >= 1:
+                    print(colored("\nREMEMBER THIS TASK-ADVICE PAIR", "light_yellow"))
+                self.memory.add([{"role": "user", "content": f"Task: {general_task}\nAdvice: {advice}"}], agent_id=self.agent_id)
+                memo_added = True
+
+        response = self._analyze(
+            comment,
+            "Does the TEXT contain information that could be committed to memory? Answer with just one word, yes or no.",
+        )
+
+        if "yes" in response.lower():
+            question = self._analyze(
+                comment,
+                "Imagine that the user forgot this information in the TEXT. How would they ask you for this information? Include no other text in your response.",
+            )
+
+            answer = self._analyze(
+                comment, "Copy the information from the TEXT that should be committed to memory. Add no explanation."
+            )
+
+            if self.verbosity >= 1:
+                print(colored("\nREMEMBER THIS QUESTION-ANSWER PAIR", "light_yellow"))
+            self.memory.add([{"role": "user", "content": f"Question: {question}\nAnswer: {answer}"}], agent_id=self.agent_id)
+            memo_added = True
+
+    def _consider_memo_retrieval(self, comment: Union[Dict, str]):
+        if self.verbosity >= 1:
+            print(colored("\nLOOK FOR RELEVANT MEMOS, AS QUESTION-ANSWER PAIRS", "light_yellow"))
+        memo_list = self._retrieve_relevant_memos(comment)
+
+        response = self._analyze(
+            comment,
+            "Does any part of the TEXT ask the agent to perform a task or solve a problem? Answer with just one word, yes or no.",
+        )
+
+        if "yes" in response.lower():
+            if self.verbosity >= 1:
+                print(colored("\nLOOK FOR RELEVANT MEMOS, AS TASK-ADVICE PAIRS", "light_yellow"))
+            task = self._analyze(
+                comment, "Copy just the task from the TEXT, then stop. Don't solve it, and don't include any advice."
+            )
+
+            general_task = self._analyze(
+                task,
+                "Summarize very briefly, in general terms, the type of task described in the TEXT. Leave out details that might not appear in a similar problem.",
+            )
+
+            memo_list.extend(self._retrieve_relevant_memos(general_task))
+
+        memo_list = list(set(memo_list))
+        return comment + self._concatenate_memo_texts(memo_list)
+
+    def _retrieve_relevant_memos(self, input_text: str) -> list:
+        search_results = self.memory.search(input_text, agent_id=self.agent_id, limit=self.max_num_retrievals)
+        memo_list = [result["memory"] for result in search_results if result["score"] <= self.recall_threshold]
+
+        if self.verbosity >= 1 and not memo_list:
+            print(colored("\nTHE CLOSEST MEMO IS BEYOND THE THRESHOLD:", "light_yellow"))
+            if search_results["results"]:
+                print(search_results["results"][0])
+            print()
+
+        return memo_list
+
+    def _concatenate_memo_texts(self, memo_list: list) -> str:
+        memo_texts = ""
+        if memo_list:
+            info = "\n# Memories that might help\n"
+            for memo in memo_list:
+                info += f"- {memo}\n"
+            if self.verbosity >= 1:
+                print(colored(f"\nMEMOS APPENDED TO LAST MESSAGE...\n{info}\n", "light_yellow"))
+            memo_texts += "\n" + info
+        return memo_texts
+
+    def _analyze(self, text_to_analyze: Union[Dict, str], analysis_instructions: Union[Dict, str]):
+        self.analyzer.reset()
+        self.teachable_agent.send(
+            recipient=self.analyzer, message=text_to_analyze, request_reply=False, silent=(self.verbosity < 2)
+        )
+        self.teachable_agent.send(
+            recipient=self.analyzer, message=analysis_instructions, request_reply=True, silent=(self.verbosity < 2)
+        )
+        return self.teachable_agent.last_message(self.analyzer)["content"]