Class Mem0Client

evaluation/scripts/utils/client.py:51–92

class Mem0Client:
    def __init__(self, enable_graph=False):
        from mem0 import MemoryClient

        self.client = MemoryClient(api_key=os.getenv("MEM0_API_KEY"))
        self.enable_graph = enable_graph

    def add(self, messages, user_id, timestamp, batch_size=2):
        max_retries = 5
        for i in range(0, len(messages), batch_size):
            batch_messages = messages[i : i + batch_size]
            for attempt in range(max_retries):
                try:
                    if self.enable_graph:
                        self.client.add(
                            messages=batch_messages,
                            timestamp=timestamp,
                            user_id=user_id,
                            enable_graph=True,
                        )
                    else:
                        self.client.add(
                            messages=batch_messages,
                            timestamp=timestamp,
                            user_id=user_id,
                        )
                    break
                except Exception as e:
                    if attempt < max_retries - 1:
                        time.sleep(2**attempt)
                    else:
                        raise e

    def search(self, query, user_id, top_k):
        res = self.client.search(
            query=query,
            top_k=top_k,
            user_id=user_id,
            enable_graph=self.enable_graph,
            filters={"AND": [{"user_id": f"{user_id}"}]},
        )
        return res


class MemobaseClient: