Class MemosApiOnlineClient

evaluation/scripts/utils/client.py:198–274

class MemosApiOnlineClient:
    def __init__(self):
        self.memos_url = os.getenv("MEMOS_ONLINE_URL")
        self.headers = {"Content-Type": "application/json", "Authorization": os.getenv("MEMOS_KEY")}

    def add(self, messages, user_id, conv_id=None, batch_size: int = 9999):
        url = f"{self.memos_url}/add/message"
        for i in range(0, len(messages), batch_size):
            batch_messages = messages[i : i + batch_size]
            payload = json.dumps(
                {
                    "messages": batch_messages,
                    "user_id": user_id,
                    "conversation_id": conv_id,
                }
            )

            max_retries = 5
            for attempt in range(max_retries):
                try:
                    response = requests.request("POST", url, data=payload, headers=self.headers)
                    assert response.status_code == 200, response.text
                    assert json.loads(response.text)["message"] == "ok", response.text
                    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):
        """Search memories."""
        url = f"{self.memos_url}/search/memory"
        payload = json.dumps(
            {
                "query": query,
                "user_id": user_id,
                "memory_limit_number": top_k,
                "mode": os.getenv("SEARCH_MODE", "fast"),
                "include_preference": True,
                "pref_top_k": 6,
            }
        )

        max_retries = 5
        for attempt in range(max_retries):
            try:
                response = requests.request("POST", url, data=payload, headers=self.headers)
                assert response.status_code == 200, response.text
                assert json.loads(response.text)["message"] == "ok", response.text
                text_mem_res = json.loads(response.text)["data"]["memory_detail_list"]
                pref_mem_res = json.loads(response.text)["data"]["preference_detail_list"]
                preference_note = json.loads(response.text)["data"]["preference_note"]
                for i in text_mem_res:
                    i.update({"memory": i.pop("memory_value")})
                explicit_pref_string = "Explicit Preference:"
                implicit_pref_string = "\n\nImplicit Preference:"
                explicit_idx = 0