From cff9c2339e1121b8695b9221088719571e2fdd1f Mon Sep 17 00:00:00 2001 From: Michael Barnathan Date: Fri, 14 Jun 2024 00:12:26 -0400 Subject: [PATCH 1/2] Initialize Client.llm_tracker on all codepaths. Prevents an issue with the check in stop_instrumenting(). To properly test this, I also needed to add a parameter to override the behavior of @singleton. --- agentops/client.py | 1 + agentops/helpers.py | 3 ++- tests/test_teardown.py | 7 ++++++- 3 files changed, 9 insertions(+), 2 deletions(-) diff --git a/agentops/client.py b/agentops/client.py index 7108e3939..8be6e13e7 100644 --- a/agentops/client.py +++ b/agentops/client.py @@ -98,6 +98,7 @@ def __init__( os.environ.get("AGENTOPS_ENV_DATA_OPT_OUT", "False").lower() == "true" ) + self.llm_tracker = None self.config = None try: diff --git a/agentops/helpers.py b/agentops/helpers.py index 24b623a3c..a63c2b1e9 100644 --- a/agentops/helpers.py +++ b/agentops/helpers.py @@ -23,7 +23,8 @@ def singleton(class_): instances = {} def getinstance(*args, **kwargs): - if class_ not in instances: + allow_multiple_instances = kwargs.pop("allow_multiple_instances", False) + if allow_multiple_instances or class_ not in instances: instances[class_] = class_(*args, **kwargs) return instances[class_] diff --git a/tests/test_teardown.py b/tests/test_teardown.py index 2a59e297f..0da22a935 100644 --- a/tests/test_teardown.py +++ b/tests/test_teardown.py @@ -20,8 +20,13 @@ class TestSessions: def setup_method(self): self.api_key = "random_api_key" self.event_type = "test_event_type" - self.client = Client(self.api_key) + self.client = Client(self.api_key, allow_multiple_instances=False) def test_exit(self): # Tests should not hang. ... + + def test_can_stop_with_no_instrumentation(self): + non_instrumented = Client(self.api_key, instrument_llm_calls=False, + allow_multiple_instances=True) + non_instrumented.stop_instrumenting() From fc0796364ca98d5b70b2730d8827ee7a6036acaf Mon Sep 17 00:00:00 2001 From: Michael Barnathan Date: Fri, 14 Jun 2024 00:22:04 -0400 Subject: [PATCH 2/2] Let's also test instrument_llm_calls=True. --- tests/test_teardown.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/test_teardown.py b/tests/test_teardown.py index 0da22a935..455d0b756 100644 --- a/tests/test_teardown.py +++ b/tests/test_teardown.py @@ -29,4 +29,10 @@ def test_exit(self): def test_can_stop_with_no_instrumentation(self): non_instrumented = Client(self.api_key, instrument_llm_calls=False, allow_multiple_instances=True) + assert non_instrumented.llm_tracker is None non_instrumented.stop_instrumenting() + + def test_initializes_llm_tracker_when_enabled(self): + instrumented = Client(self.api_key, instrument_llm_calls=True, + allow_multiple_instances=True) + assert instrumented.llm_tracker is not None