From 767b4d9e2b0559fd528ff232e17e176714ee55dd Mon Sep 17 00:00:00 2001
From: Alex Reibman <areibman@gmail.com>
Date: Wed, 1 Nov 2023 23:30:40 -0700
Subject: [PATCH] 26 tracebacks for failed events (#43)

* added failstrings to events

* added api key checker

* update tox

* added daemon

* revert test

* added fail agent for mock testing

* fix typo
---
 agentops/client.py  | 16 ++++++++++++----
 agentops/worker.py  |  1 +
 tests/fail_agent.py | 39 +++++++++++++++++++++++++++++++++++++++
 tests/mock_agent.py |  4 ++--
 tox.ini             |  1 +
 5 files changed, 55 insertions(+), 6 deletions(-)
 create mode 100644 tests/fail_agent.py

diff --git a/agentops/client.py b/agentops/client.py
index cec565b6..7107c257 100644
--- a/agentops/client.py
+++ b/agentops/client.py
@@ -12,6 +12,7 @@
 from uuid import uuid4
 from typing import Optional, List
 from pydantic import Field
+from os import environ
 import functools
 import logging
 import inspect
@@ -28,7 +29,7 @@ class Client:
     Client for AgentOps service.
 
     Args:
-        api_key (str): API Key for AgentOps services.
+        api_key (str, optional): API Key for AgentOps services. If none is provided, key will be read from the AGENTOPS_API_KEY environment variable.
         tags (List[str], optional): Tags for the sessions that can be used for grouping or sorting later (e.g. ["GPT-4"]).
         endpoint (str, optional): The endpoint for the AgentOps service. Defaults to 'https://agentops-server-v2.fly.dev'.
         max_wait_time (int, optional): The maximum time to wait in milliseconds before flushing the queue. Defaults to 1000.
@@ -37,11 +38,18 @@ class Client:
         session (Session, optional): A Session is a grouping of events (e.g. a run of your agent).
     """
 
-    def __init__(self, api_key: str, tags: Optional[List[str]] = None,
+    def __init__(self, api_key: Optional[str] = None, tags: Optional[List[str]] = None,
                  endpoint: Optional[str] = 'https://agentops-server-v2.fly.dev',
                  max_wait_time: Optional[int] = 1000,
                  max_queue_size: Optional[int] = 100):
 
+        # Get API key from env
+        if api_key is None:
+            api_key = environ.get('AGENTOPS_API_KEY')
+
+        if api_key is None:
+            print("AgentOps API key not provided. Session data will not be recorded.")
+
         # Create a worker config
         self.config = Configuration(api_key, endpoint,
                                     max_wait_time, max_queue_size)
@@ -163,7 +171,7 @@ def _record_event_sync(self, func, event_name, tags, *args, **kwargs):
             # Record the event after the function call
             self.record(Event(event_type=event_name,
                               params=arg_values,
-                              returns=None,
+                              returns=str(e),
                               result='Fail',
                               action_type='action',
                               init_timestamp=init_time,
@@ -207,7 +215,7 @@ async def _record_event_async(self, func, event_name, tags, *args, **kwargs):
             # Record the event after the function call
             self.record(Event(event_type=event_name,
                               params=arg_values,
-                              returns=None,
+                              returns=str(e),
                               result='Fail',
                               action_type='action',
                               init_timestamp=init_time,
diff --git a/agentops/worker.py b/agentops/worker.py
index 75154821..f41d01d2 100644
--- a/agentops/worker.py
+++ b/agentops/worker.py
@@ -39,6 +39,7 @@ def __init__(self, config: Configuration) -> None:
         self.lock = threading.Lock()
         self.stop_flag = threading.Event()
         self.thread = threading.Thread(target=self.run)
+        self.thread.daemon = True
         self.thread.start()
 
     def add_event(self, event: dict) -> None:
diff --git a/tests/fail_agent.py b/tests/fail_agent.py
new file mode 100644
index 00000000..ed18543b
--- /dev/null
+++ b/tests/fail_agent.py
@@ -0,0 +1,39 @@
+import agentops
+import time
+import asyncio
+from dotenv import load_dotenv
+
+load_dotenv()
+
+
+print('init')
+
+ao_client = agentops.Client(tags=['mock tests'])
+
+
+@ao_client.record_action('fail')
+def fail_func(sleep):
+    time.sleep(sleep)
+    print('sync sleep')
+
+    print(1/0)
+
+
+@ao_client.record_action('sleep')
+def sleep_func(sleep):
+    time.sleep(sleep)
+    print('sync sleep')
+
+
+async def main():
+
+    print('Action 1')
+    sleep_func(0.1)
+    fail_func(0.1)
+
+    ao_client.end_session(end_state='Success')
+
+
+if __name__ == '__main__':
+    asyncio.run(main())
+    print('done')
diff --git a/tests/mock_agent.py b/tests/mock_agent.py
index 50a046a8..d7fcee75 100644
--- a/tests/mock_agent.py
+++ b/tests/mock_agent.py
@@ -9,8 +9,7 @@
 
 print('init')
 
-ao_client = agentops.Client(api_key=os.environ['AGENTOPS_API_KEY'],
-                            tags=['mock tests'])
+ao_client = agentops.Client(tags=['mock tests'])
 
 
 @ao_client.record_action('action')
@@ -68,3 +67,4 @@ async def main():
 
 if __name__ == '__main__':
     asyncio.run(main())
+    print('done')
diff --git a/tox.ini b/tox.ini
index 25fac5bf..d1420bc4 100644
--- a/tox.ini
+++ b/tox.ini
@@ -9,6 +9,7 @@ envlist = py37, py38, py39, mypy
 [testenv]
 deps = 
     pytest
+    pytest-asyncio
     requests_mock
     coverage
     pydantic