add stop_instrumenting

agentops/__init__.py

@@ -9,6 +9,7 @@
 from .decorators import record_function
 from .agent import track_agent
 from .log_config import set_logging_level_info, set_logging_level_critial
+from .langchain_callback_handler import LangchainCallbackHandler, AsyncLangchainCallbackHandler
 
 
 def init(api_key: Optional[str] = None,
@@ -128,3 +129,6 @@ def set_parent_key(parent_key):
             parent_key (str): The API key of the parent organization to set.
     """
     Client().set_parent_key(parent_key)
+
+def stop_instrumenting():
+    Client().stop_instrumenting()

agentops/client.py

@@ -136,7 +136,7 @@ def record(self, event: Event | ErrorEvent):
             Args:
                 event (Event): The event to record.
         """
-        if not event.end_timestamp or event.init_timestamp == event.end_timestamp:
+        if isinstance(event, Event) and not event.end_timestamp or event.init_timestamp == event.end_timestamp:
             event.end_timestamp = get_ISO_time()
         if self._session is not None and not self._session.has_ended:
             if isinstance(event, ErrorEvent):
@@ -356,3 +356,6 @@ def set_parent_key(self, parent_key: str):
     @property
     def parent_key(self):
         return self.config.parent_key
+
+    def stop_instrumenting(self):
+        self.llm_tracker.stop_instrumenting()

agentops/llm_tracker.py

@@ -9,6 +9,8 @@
 import inspect
 import pprint
 
+original_create = None
+original_create_async = None
 
 class LlmTracker:
     SUPPORTED_APIS = {
@@ -227,6 +229,7 @@ def override_openai_v1_completion(self):
         from openai.resources.chat import completions
 
         # Store the original method
+        global original_create
         original_create = completions.Completions.create
 
         def patched_function(*args, **kwargs):
@@ -242,12 +245,13 @@ def override_openai_v1_async_completion(self):
         from openai.resources.chat import completions
 
         # Store the original method
-        original_create = completions.AsyncCompletions.create
+        global original_create_async
+        original_create_async = completions.AsyncCompletions.create
 
         async def patched_function(*args, **kwargs):
             # Call the original function with its original arguments
             init_timestamp = get_ISO_time()
-            result = await original_create(*args, **kwargs)
+            result = await original_create_async(*args, **kwargs)
             return self._handle_response_v1_openai(result, kwargs, init_timestamp)
 
         # Override the original method with the patched one
@@ -342,3 +346,17 @@ def override_api(self):
                             # Patch openai <v1.0.0 methods
                             for method_path in self.SUPPORTED_APIS['openai']['0.0.0']:
                                 self._override_method(api, method_path, module)
+
+    def stop_instrumenting(self):
+        self.undo_override_openai_v1_async_completion()
+        self.undo_override_openai_v1_completion()
+
+    def undo_override_openai_v1_completion(self):
+        global original_create
+        from openai.resources.chat import completions
+        completions.Completions.create = original_create
+
+    def undo_override_openai_v1_async_completion(self):
+        global original_create_async
+        from openai.resources.chat import completions
+        original_create_async = completions.AsyncCompletions.create

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "agentops"
-version = "0.1.7"
+version = "0.1.8"
 authors = [
   { name="Alex Reibman", email="areibman@gmail.com" },
   { name="Shawn Qiu", email="siyangqiu@gmail.com" },