From ecea9ab757d34e053e83be83c9ff605b2ee4956f Mon Sep 17 00:00:00 2001
From: Teo <teocns@gmail.com>
Date: Thu, 16 Jan 2025 20:17:43 +0100
Subject: [PATCH] Remove agentops.enums module; move EndState to .session,
 EventType to .event

Signed-off-by: Teo <teocns@gmail.com>
---
 agentops/enums.py                   | 26 --------------------------
 agentops/event.py                   | 12 ++++++++++--
 agentops/partners/autogen_logger.py |  2 +-
 agentops/session.py                 | 28 ++++++++++++++++++++--------
 4 files changed, 31 insertions(+), 37 deletions(-)
 delete mode 100644 agentops/enums.py

diff --git a/agentops/enums.py b/agentops/enums.py
deleted file mode 100644
index cce033921..000000000
--- a/agentops/enums.py
+++ /dev/null
@@ -1,26 +0,0 @@
-from enum import Enum
-
-
-class EventType(Enum):
-    LLM = "llms"
-    ACTION = "actions"
-    API = "apis"
-    TOOL = "tools"
-    ERROR = "errors"
-
-
-class EndState(Enum):
-    """
-    Enum representing the possible end states of a session.
-
-    Attributes:
-        SUCCESS: Indicates the session ended successfully.
-        FAIL: Indicates the session failed.
-        INDETERMINATE (default): Indicates the session ended with an indeterminate state.
-                       This is the default state if not specified, e.g. if you forget to call end_session()
-                       at the end of your program or don't pass it the end_state parameter
-    """
-
-    SUCCESS = "Success"
-    FAIL = "Fail"
-    INDETERMINATE = "Indeterminate"  # Default
diff --git a/agentops/event.py b/agentops/event.py
index 1ed356bfb..abffcacc4 100644
--- a/agentops/event.py
+++ b/agentops/event.py
@@ -7,19 +7,27 @@
 
 import traceback
 from dataclasses import dataclass, field
+from enum import Enum
 from typing import Any, Dict, List, Optional, Sequence, Union
 from uuid import UUID, uuid4
 
-from .enums import EventType
 from .helpers import check_call_stack_for_agent_id, get_ISO_time
 
 
+class EventType(Enum):
+    LLM = "llms"
+    ACTION = "actions"
+    API = "apis"
+    TOOL = "tools"
+    ERROR = "errors"
+
+
 @dataclass
 class Event:
     """
     Abstract base class for events that will be recorded. Should not be instantiated directly.
 
-    event_type(str): The type of event. Defined in enums.EventType. Some values are 'llm', 'action', 'api', 'tool', 'error'.
+    event_type(str): The type of event. Defined in events.EventType. Some values are 'llm', 'action', 'api', 'tool', 'error'.
     params(dict, optional): The parameters of the function containing the triggered event, e.g. {'x': 1} in example below
     returns(str, optional): The return value of the function containing the triggered event, e.g. 2 in example below
     init_timestamp(str): A timestamp indicating when the event began. Defaults to the time when this Event was instantiated.
diff --git a/agentops/partners/autogen_logger.py b/agentops/partners/autogen_logger.py
index e35b04c80..9fc85fb28 100644
--- a/agentops/partners/autogen_logger.py
+++ b/agentops/partners/autogen_logger.py
@@ -11,7 +11,7 @@
 
 from autogen.logger.base_logger import BaseLogger, LLMConfig
 
-from agentops.enums import EndState
+from agentops.session import EndState
 from agentops.helpers import get_ISO_time
 
 from agentops import LLMEvent, ToolEvent, ActionEvent
diff --git a/agentops/session.py b/agentops/session.py
index b9f07d20b..95d1fba15 100644
--- a/agentops/session.py
+++ b/agentops/session.py
@@ -6,6 +6,7 @@
 import threading
 from datetime import datetime, timezone
 from decimal import ROUND_HALF_UP, Decimal
+from enum import Enum
 from typing import Any, Dict, List, Optional, Sequence, Union
 from uuid import UUID, uuid4
 
@@ -13,16 +14,10 @@
 from opentelemetry.context import attach, detach, set_value
 from opentelemetry.sdk.resources import SERVICE_NAME, Resource
 from opentelemetry.sdk.trace import ReadableSpan, TracerProvider
-from opentelemetry.sdk.trace.export import (
-    BatchSpanProcessor,
-    ConsoleSpanExporter,
-    SpanExporter,
-    SpanExportResult,
-)
+from opentelemetry.sdk.trace.export import BatchSpanProcessor, ConsoleSpanExporter, SpanExporter, SpanExportResult
 from termcolor import colored
 
 from .config import Configuration
-from .enums import EndState
 from .event import ErrorEvent, Event
 from .exceptions import ApiServerException
 from .helpers import filter_unjsonable, get_ISO_time, safe_serialize
@@ -63,6 +58,23 @@
 """
 
 
+class EndState(Enum):
+    """
+    Enum representing the possible end states of a session.
+
+    Attributes:
+        SUCCESS: Indicates the session ended successfully.
+        FAIL: Indicates the session failed.
+        INDETERMINATE (default): Indicates the session ended with an indeterminate state.
+                       This is the default state if not specified, e.g. if you forget to call end_session()
+                       at the end of your program or don't pass it the end_state parameter
+    """
+
+    SUCCESS = "Success"
+    FAIL = "Fail"
+    INDETERMINATE = "Indeterminate"  # Default
+
+
 class SessionExporter(SpanExporter):
     """
     Manages publishing events for Session
@@ -290,7 +302,7 @@ def end_session(
                 return None
 
             if not any(end_state == state.value for state in EndState):
-                logger.warning("Invalid end_state. Please use one of the EndState enums")
+                logger.warning("Invalid end_state. Please use one of the EndState")
                 return None
 
             try: