1 create enum for session end states and event results

agentops/agentops.py

@@ -6,10 +6,9 @@
 """
 
 from .config import Configuration
-from .event import Session, Event
+from .event import Session, Event, EventState
 from .worker import Worker
 from uuid import uuid4
-import json
 from typing import Optional, Dict
 import functools
 import inspect
@@ -63,29 +62,28 @@ def wrapper(*args, **kwargs):
                 arg_names = list(func_args.keys())
                 arg_values = dict(zip(arg_names, args))
                 arg_values.update(kwargs)
-
                 try:
-                    output = func(*args, **kwargs)
+                    returns = func(*args, **kwargs)
 
                     # Record the event after the function call
                     self.record(Event(event_type=event_name,
                                       params=arg_values,
-                                      output=output,
-                                      result="SUCCESS",
+                                      returns=returns,
+                                      result="Success",
                                       tags=tags))
 
                 except Exception as e:
                     # Record the event after the function call
                     self.record(Event(event_type=event_name,
                                       params=arg_values,
-                                      output=None,
-                                      result='FAIL',
+                                      returns=None,
+                                      result='Fail',
                                       tags=tags))
 
                     # Re-raise the exception
                     raise
 
-                return output
+                return returns
 
             return wrapper
 
@@ -102,13 +100,18 @@ def start_session(self, tags: Optional[Dict[str, str]] = None):
         self.worker = Worker(self.config)
         self.worker.start_session(self.session)
 
-    def end_session(self, end_state: Optional[str] = None, rating: Optional[str] = None):
+    def end_session(self, end_state: EventState = EventState.INDETERMINATE, rating: Optional[str] = None):
         """
         End the current session with the AgentOps service.
 
         Args:
             end_state (str, optional): The final state of the session.
             rating (str, optional): The rating for the session.
         """
+        valid_results = set(vars(EventState).values())
+        if end_state not in valid_results:
+            raise ValueError(
+                f"end_state must be one of {EventState.__args__}. Provided: {end_state}")
+
         self.session.end_session(end_state, rating)
         self.worker.end_session(self.session)

agentops/event.py

@@ -1,13 +1,24 @@
-from .helpers import get_ISO_time
-from typing import Optional, Dict, Literal
-
 """
 AgentOps events.
 
 Classes:
     Event: Represents discrete events to be recorded.
     Session: Represents a session of events, with a start and end state.
 """
+from .helpers import get_ISO_time
+from typing import Optional, Dict
+
+
+class EventState:
+    SUCCESS = "Success"
+    FAIL = "Fail"
+    INDETERMINATE = "Indeterminate"
+
+
+class SessionState:
+    SUCCESS = "Success"
+    FAIL = "Fail"
+    INDETERMINATE = "Indeterminate"
 
 
 class Event:
@@ -17,8 +28,8 @@ class Event:
     Args:
         event_type (str): Type of the event, e.g., "API Call". Required.
         params (str, optional): The parameters passed to the operation.
-        output (str, optional): The output of the operation.
-        result (str, optional): Result of the operation, e.g., "success", "fail", "indeterminate".
+        returns (str, optional): The output of the operation.
+        result (str, optional): Result of the operation, e.g., "Success", "Fail", "Indeterminate".
         tags (Dict[str, str], optional): Tags that can be used for grouping or sorting later. e.g. {"llm": "GPT-4"}.
 
 
@@ -28,13 +39,13 @@ class Event:
 
     def __init__(self, event_type: str,
                  params: Optional[str] = None,
-                 output: Optional[str] = None,
-                 result: Optional[str] = None,
+                 returns: Optional[str] = None,
+                 result: EventState = EventState.INDETERMINATE,
                  tags: Optional[Dict[str, str]] = None
                  ):
         self.event_type = event_type
         self.params = params
-        self.output = output
+        self.returns = returns
         self.result = result
         self.tags = tags
         self.timestamp = get_ISO_time()
@@ -61,14 +72,18 @@ def __init__(self, session_id: str, tags: Optional[Dict[str, str]] = None):
         self.init_timestamp = get_ISO_time()
         self.tags = tags
 
-    def end_session(self, end_state: Optional[str], rating: Optional[str] = None):
+    def end_session(self, end_state: SessionState = SessionState.INDETERMINATE, rating: Optional[str] = None):
         """
         End the session with a specified state and rating.
 
         Args:
             end_state (str, optional): The final state of the session. Suggested: "Success", "Fail", "Indeterminate"
             rating (str, optional): The rating for the session.
         """
+        valid_results = set(vars(SessionState).values())
+        if end_state not in valid_results:
+            raise ValueError(
+                f"end_state must be one of {valid_results}. Provided: {end_state}")
         self.end_state = end_state
         self.rating = rating
         self.end_timestamp = get_ISO_time()

agentops/http.py

@@ -69,7 +69,7 @@ def post(url: str, payload: bytes, api_key: str = None, header=None) -> Response
                 JSON_HEADER["X-Agentops-Auth"] = api_key
 
             res = request_session.post(url, data=payload,
-                                           headers=JSON_HEADER, timeout=20)
+                                       headers=JSON_HEADER, timeout=20)
 
             result.parse(res)
         except requests.exceptions.Timeout:

agentops/logger.py

@@ -1,26 +1,27 @@
 import logging
 import re
 from .agentops import AgentOps
-from .event import Event
+from .event import Event, EventState
+
 
 class AgentOpsLogger():
     """
     A utility class for creating loggers and handlers configured to work with the AgentOps service.
-    
+
     This class provides two static methods for creating a logger or a handler that sends log 
     records to the AgentOps service. The logger and handler are configured with a specific 
     AgentOps client and name.
-    
+
     Example Usage:
-    
+
     >>> from agentops import AgentOps
     >>> client = AgentOps(...)
     >>> logger = AgentOpsLogger.get_agentops_logger(client, 'my_logger')
     >>> logger.info('This is an info log')
-    
+
     This will send an 'info' log to the AgentOps service.
     """
-    
+
     @staticmethod
     def get_agentops_logger(client: AgentOps, name: str, level=logging.DEBUG):
         """
@@ -40,7 +41,7 @@ def get_agentops_logger(client: AgentOps, name: str, level=logging.DEBUG):
         handler.setLevel(level)
         logger.addHandler(handler)
         return logger
-        
+
     @staticmethod
     def get_agentops_handler(client: AgentOps, name: str):
         """
@@ -103,9 +104,9 @@ def emit(self, record):
         log_entry = self.remove_color_codes(log_entry)
 
         if record.levelno == logging.ERROR:
-            result = "fail"
+            result = EventState.FAIL
         else:
-            result = "indeterminate"
+            result = EventState.INDETERMINATE
 
         self.client.record(
-            Event(f'{self.name}:{record.levelname}', output=log_entry, result=result))
+            Event(f'{self.name}:{record.levelname}', returns=log_entry, result=result))

tests/test_canary.py

@@ -3,6 +3,7 @@
 import time
 
 from agentops import AgentOps, Event, Configuration
+from agentops.event import EventState
 
 
 @pytest.fixture
@@ -13,6 +14,7 @@ def mock_req():
         m.post(url + '/sessions', text='ok')
         yield m
 
+
 class TestCanary:
     def setup_method(self):
         self.url = 'https://agentops-server-v2.fly.dev'
@@ -21,7 +23,7 @@ def setup_method(self):
         self.client = AgentOps(api_key=self.api_key, config=self.config)
 
     def teardown_method(self):
-        self.client.end_session(end_state="success")
+        self.client.end_session(end_state=EventState.SUCCESS)
 
     def test_agent_ops_record(self, mock_req):
         # Arrange

tests/test_logger.py

@@ -2,6 +2,7 @@
 import requests_mock
 import time
 from agentops import AgentOps, AgentOpsLogger, Configuration
+from agentops.event import EventState
 
 
 @pytest.fixture
@@ -20,9 +21,8 @@ def setup_method(self):
         self.config = Configuration(max_wait_time=5)
         self.client = AgentOps(api_key=self.api_key, config=self.config)
 
-
     def teardown_method(self):
-        self.client.end_session(end_state="success")
+        self.client.end_session(end_state=EventState.SUCCESS)
 
     def test_info(self, mock_req):
         # Arrange
@@ -37,15 +37,14 @@ def test_info(self, mock_req):
         except Exception as e:
             pytest.fail(f"test_info failed with {e}")
 
-
         time.sleep(0.1)
 
         # Assert
         assert len(mock_req.request_history) == 1
         assert mock_req.last_request.headers['X-Agentops-Auth'] == self.api_key
         request_json = mock_req.last_request.json()
         assert request_json['events'][0]['event_type'] == f"{self.event_type}:INFO"
-        assert request_json['events'][0]['output'] == test_message
+        assert request_json['events'][0]['returns'] == test_message
 
     def test_error(self, mock_req):
         # Arrange
@@ -67,7 +66,7 @@ def test_error(self, mock_req):
         assert mock_req.last_request.headers['X-Agentops-Auth'] == self.api_key
         request_json = mock_req.last_request.json()
         assert request_json['events'][0]['event_type'] == f"{self.event_type}:ERROR"
-        assert request_json['events'][0]['output'] == test_message
+        assert request_json['events'][0]['returns'] == test_message
 
     def test_warn(self, mock_req):
         # Arrange
@@ -89,4 +88,4 @@ def test_warn(self, mock_req):
         assert mock_req.last_request.headers['X-Agentops-Auth'] == self.api_key
         request_json = mock_req.last_request.json()
         assert request_json['events'][0]['event_type'] == f"{self.event_type}:WARNING"
-        assert request_json['events'][0]['output'] == test_message
+        assert request_json['events'][0]['returns'] == test_message

tests/test_session.py

@@ -3,6 +3,7 @@
 import time
 
 from agentops import AgentOps, Event, Configuration
+from agentops.event import EventState
 
 
 @pytest.fixture
@@ -27,25 +28,26 @@ def test_session(self, mock_req):
         # Act
         client.record(Event(self.event_type))
 
-        # Assert
-        assert len(mock_req.request_history) == 0
+        # Assert the session has been initiated and the id has been created on backend.
+        assert len(mock_req.request_history) == 1
 
         # Act
         client.record(Event(self.event_type))
         time.sleep(0.1)
 
-        # Assert
-        assert len(mock_req.request_history) == 1
+        # Assert an event has been added
+        assert len(mock_req.request_history) == 2
         assert mock_req.last_request.headers['X-Agentops-Auth'] == self.api_key
         request_json = mock_req.last_request.json()
         assert request_json['events'][0]['event_type'] == self.event_type
 
         # Act
-        end_state = "Succeed"
+        end_state = EventState.SUCCESS
         client.end_session(end_state)
         time.sleep(0.1)
 
-        assert len(mock_req.request_history) == 2
+        # Since a session has ended, no more events should be recorded, but end_session should be called
+        assert len(mock_req.request_history) == 3
         assert mock_req.last_request.headers['X-Agentops-Auth'] == self.api_key
         request_json = mock_req.last_request.json()
         assert request_json['session']['rating'] == None
@@ -62,22 +64,23 @@ def test_tags(self, mock_req):
         client.record(Event(self.event_type))
         time.sleep(0.1)
 
-        # Assert
-        assert len(mock_req.request_history) == 1
+        # Assert 2 requests - 1 for session init, 1 for event
+        assert len(mock_req.request_history) == 2
         assert mock_req.last_request.headers['X-Agentops-Auth'] == self.api_key
         request_json = mock_req.last_request.json()
         assert request_json['events'][0]['event_type'] == self.event_type
 
         # Act
-        end_state = "Succeed"
-        client.end_session()
+        end_state = EventState.SUCCESS
+        client.end_session(end_state)
         time.sleep(0.1)
 
-        assert len(mock_req.request_history) == 2
+        # Assert 3 requets, 1 for session init, 1 for event, 1 for end session
+        assert len(mock_req.request_history) == 3
         assert mock_req.last_request.headers['X-Agentops-Auth'] == self.api_key
         request_json = mock_req.last_request.json()
         assert request_json['session']['rating'] == None
-        assert request_json['session']['end_state'] == None
+        assert request_json['session']['end_state'] == end_state
         assert request_json['session']['tags'] == tags
 
 
@@ -90,7 +93,7 @@ def setup_method(self):
         self.client = AgentOps(self.api_key, config=self.config)
 
     def teardown_method(self):
-        self.client.end_session(end_state="success")
+        self.client.end_session(end_state=EventState.SUCCESS)
 
     def test_record_action_decorator(self, mock_req):
         @self.client.record_action(event_name=self.event_type, tags={'foo': 'bar'})
@@ -107,8 +110,8 @@ def dummy_func(x, y):
         assert request_json['event']['event_type'] == self.event_type
         assert request_json['event']['params'] == {
             'args': [3, 4], 'kwargs': {}}
-        assert request_json['event']['output'] == 7
-        assert request_json['event']['result'] == 'SUCCESS'
+        assert request_json['event']['returns'] == 7
+        assert request_json['event']['result'] == EventState.SUCCESS
         assert request_json['event']['tags'] == {'foo': 'bar'}
 
     def test_record_action_decorator(self, mock_req):
@@ -127,6 +130,6 @@ def dummy_func(x, y):
         request_json = mock_req.last_request.json()
         assert request_json['events'][0]['event_type'] == self.event_type
         assert request_json['events'][0]['params'] == {'x': 3, 'y': 4}
-        assert request_json['events'][0]['output'] == 7
-        assert request_json['events'][0]['result'] == 'SUCCESS'
+        assert request_json['events'][0]['returns'] == 7
+        assert request_json['events'][0]['result'] == EventState.SUCCESS
         assert request_json['events'][0]['tags'] == {'foo': 'bar'}