chore: import upstream snapshot with attribution

zep-memory-assistant/agent.py

@@ -0,0 +1,86 @@
+from typing import Union, Dict
+from autogen import ConversableAgent, Agent
+from zep_cloud.client import Zep
+from zep_cloud import Message, Memory
+
+
+class ZepConversableAgent(ConversableAgent):  # Agent with Zep memory
+    """A custom ConversableAgent that integrates with Zep for long-term memory."""
+
+    def __init__(
+        self,
+        name: str,
+        system_message: str,
+        llm_config: dict,
+        function_map: dict,
+        human_input_mode: str,
+        zep_session_id: str,
+        zep_client: Zep,
+        min_fact_rating: float,
+    ):
+        super().__init__(
+            name=name,
+            system_message=system_message,
+            llm_config=llm_config,
+            human_input_mode=human_input_mode,
+            function_map=function_map,
+        )
+        self.zep_session_id = zep_session_id
+        self.zep_client = zep_client
+        self.min_fact_rating = min_fact_rating
+        # Store the original system message as we will update it with relevant facts from Zep
+        self.original_system_message = system_message
+        self.register_hook(
+            "process_message_before_send", self._zep_persist_assistant_messages
+        )
+        # Note: Persisting user messages needs to happen *before* the agent
+        # processes them to fetch relevant facts. We'll handle this outside
+        # the hook based on Streamlit input.
+
+    def _zep_persist_assistant_messages(
+        self,
+        message: Union[Dict, str],
+        sender: Agent,
+        recipient: Agent,
+        silent: bool,
+    ):
+        """Agent sends a message to the user. Add the message to Zep."""
+        if sender == self:
+            if isinstance(message, dict):
+                content = message.get("content", "")
+            else:
+                content = str(message)
+
+            if content:
+                zep_message = Message(
+                    role_type="assistant", role=self.name, content=content
+                )
+                self.zep_client.memory.add(
+                    session_id=self.zep_session_id, messages=[zep_message]
+                )
+        return message
+
+    def _zep_fetch_and_update_system_message(self):
+        """Fetch facts and update system message."""
+        memory: Memory = self.zep_client.memory.get(
+            self.zep_session_id, min_rating=self.min_fact_rating
+        )
+        context = memory.context or "No specific facts recalled."
+
+        # Update the system message for the next inference
+        self.update_system_message(
+            self.original_system_message
+            + f"\n\nRelevant facts about the user and prior conversation:\n{context}"
+        )
+
+    def _zep_persist_user_message(self, user_content: str, user_name: str = "User"):
+        """User sends a message to the agent. Add the message to Zep."""
+        if user_content:
+            zep_message = Message(
+                role_type="user",
+                role=user_name,
+                content=user_content,
+            )
+            self.zep_client.memory.add(
+                session_id=self.zep_session_id, messages=[zep_message]
+            )