diff --git a/mem0/llms/openai_structured.py b/mem0/llms/openai_structured.py index ab56b85f..0625c1e8 100644 --- a/mem0/llms/openai_structured.py +++ b/mem0/llms/openai_structured.py @@ -1,5 +1,6 @@ -import json import os +import json + from typing import Dict, List, Optional from openai import OpenAI diff --git a/mem0/memory/telemetry.py b/mem0/memory/telemetry.py index ec819e93..9b78d775 100644 --- a/mem0/memory/telemetry.py +++ b/mem0/memory/telemetry.py @@ -1,11 +1,20 @@ import logging import platform import sys +import os from posthog import Posthog from mem0.memory.setup import get_user_id, setup_config +MEM0_TELEMETRY = os.environ.get("MEM0_TELEMETRY", "True") + +if isinstance(MEM0_TELEMETRY, str): + MEM0_TELEMETRY = MEM0_TELEMETRY.lower() in ("true", "1", "yes") + +if not isinstance(MEM0_TELEMETRY, bool): + raise ValueError("MEM0_TELEMETRY must be a boolean value.") + logging.getLogger('posthog').setLevel(logging.CRITICAL + 1) logging.getLogger('urllib3').setLevel(logging.CRITICAL + 1) @@ -15,6 +24,9 @@ class AnonymousTelemetry: # Call setup config to ensure that the user_id is generated setup_config() self.user_id = get_user_id() + # Optional + if not MEM0_TELEMETRY: + self.posthog.disabled = True def capture_event(self, event_name, properties=None): if properties is None: @@ -64,6 +76,7 @@ def capture_event(event_name, memory_instance, additional_data=None): telemetry.capture_event(event_name, event_data) + def capture_client_event(event_name, instance, additional_data=None): event_data = { "function": f"{instance.__class__.__module__}.{instance.__class__.__name__}", diff --git a/tests/test_telemetry.py b/tests/test_telemetry.py new file mode 100644 index 00000000..aa36f2f0 --- /dev/null +++ b/tests/test_telemetry.py @@ -0,0 +1,29 @@ +import os +import pytest +from unittest.mock import patch + +MEM0_TELEMETRY = os.environ.get("MEM0_TELEMETRY", "True") + +if isinstance(MEM0_TELEMETRY, str): + MEM0_TELEMETRY = MEM0_TELEMETRY.lower() in ("true", "1", "yes") + +def use_telemetry(): + if os.getenv('MEM0_TELEMETRY', "true").lower() == "true": + return True + return False + +@pytest.fixture(autouse=True) +def reset_env(): + with patch.dict(os.environ, {}, clear=True): + yield + +def test_telemetry_enabled(): + with patch.dict(os.environ, {'MEM0_TELEMETRY': "true"}): + assert use_telemetry() is True + +def test_telemetry_disabled(): + with patch.dict(os.environ, {'MEM0_TELEMETRY': "false"}): + assert use_telemetry() is False + +def test_telemetry_default_enabled(): + assert use_telemetry() is True