Dev Khant
77 lines
2.5 KiB
Python
77 lines
2.5 KiB
Python
import argparse
|
|
import json
|
|
import os
|
|
|
|
from dotenv import load_dotenv
|
|
from tqdm import tqdm
|
|
from zep_cloud import Message
|
|
from zep_cloud.client import Zep
|
|
|
|
load_dotenv()
|
|
|
|
|
|
class ZepAdd:
|
|
def __init__(self, data_path=None):
|
|
self.zep_client = Zep(api_key=os.getenv("ZEP_API_KEY"))
|
|
self.data_path = data_path
|
|
self.data = None
|
|
if data_path:
|
|
self.load_data()
|
|
|
|
def load_data(self):
|
|
with open(self.data_path, "r") as f:
|
|
self.data = json.load(f)
|
|
return self.data
|
|
|
|
def process_conversation(self, run_id, item, idx):
|
|
conversation = item["conversation"]
|
|
|
|
user_id = f"run_id_{run_id}_experiment_user_{idx}"
|
|
session_id = f"run_id_{run_id}_experiment_session_{idx}"
|
|
|
|
# # delete all memories for the two users
|
|
# self.zep_client.user.delete(user_id=user_id)
|
|
# self.zep_client.memory.delete(session_id=session_id)
|
|
|
|
self.zep_client.user.add(user_id=user_id)
|
|
self.zep_client.memory.add_session(
|
|
user_id=user_id,
|
|
session_id=session_id,
|
|
)
|
|
|
|
print("Starting to add memories... for user", user_id)
|
|
for key in tqdm(conversation.keys(), desc=f"Processing user {user_id}"):
|
|
if key in ["speaker_a", "speaker_b"] or "date" in key:
|
|
continue
|
|
|
|
date_time_key = key + "_date_time"
|
|
timestamp = conversation[date_time_key]
|
|
chats = conversation[key]
|
|
|
|
for chat in tqdm(chats, desc=f"Adding chats for {key}", leave=False):
|
|
self.zep_client.memory.add(
|
|
session_id=session_id,
|
|
messages=[
|
|
Message(
|
|
role=chat["speaker"],
|
|
role_type="user",
|
|
content=f"{timestamp}: {chat['text']}",
|
|
)
|
|
],
|
|
)
|
|
|
|
def process_all_conversations(self, run_id):
|
|
if not self.data:
|
|
raise ValueError("No data loaded. Please set data_path and call load_data() first.")
|
|
for idx, item in tqdm(enumerate(self.data)):
|
|
if idx == 0:
|
|
self.process_conversation(run_id, item, idx)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--run_id", type=str, required=True)
|
|
args = parser.parse_args()
|
|
zep_add = ZepAdd(data_path="../../dataset/locomo10.json")
|
|
zep_add.process_all_conversations(args.run_id)
|