wehub-resource-sync
87 lines
3.1 KiB
Python
87 lines
3.1 KiB
Python
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]
|
|
)
|