add stop_instrumenting (#185)

* add stop_instrumenting

* version bump

* test deps

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: Union[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 and self._worker is not None:
             if isinstance(event, ErrorEvent):
@@ -359,3 +359,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

@@ -10,6 +10,8 @@
 from typing import Optional
 import pprint
 
+original_create = None
+original_create_async = None
 
 class LlmTracker:
     SUPPORTED_APIS = {
@@ -230,6 +232,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):
@@ -245,12 +248,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
@@ -345,3 +349,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="[email protected]" },
   { name="Shawn Qiu", email="[email protected]" },

tox.ini

@@ -17,6 +17,8 @@ deps =
     types-requests
     psutil
     openai
+    langchain-core
+    langchain
 commands =
     coverage run --source agentops -m pytest
     coverage report -m