From 3ce744dd76c44f9b066c0ece6220cca59a8c2636 Mon Sep 17 00:00:00 2001
From: Albert Kim Jr <albertkimjunior@gmail.com>
Date: Wed, 26 Jun 2024 17:10:54 -0700
Subject: [PATCH] Fixes #254 (#269)

* Fixes 254

* Deleted stale code

* Plumbing Models

* Fixed breaking change lol

* Black reformatting
---
 agentops/__init__.py |  3 +--
 agentops/enums.py    | 14 --------------
 agentops/event.py    |  8 ++++----
 3 files changed, 5 insertions(+), 20 deletions(-)

diff --git a/agentops/__init__.py b/agentops/__init__.py
index 8f8fdc9b..d3b3b7f3 100755
--- a/agentops/__init__.py
+++ b/agentops/__init__.py
@@ -6,7 +6,6 @@
 from .client import Client
 from .config import Configuration
 from .event import Event, ActionEvent, LLMEvent, ToolEvent, ErrorEvent
-from .enums import Models
 from .decorators import record_function
 from .agent import track_agent
 from .log_config import logger
@@ -118,7 +117,7 @@ def end_session(
         video (str, optional): URL to a video recording of the session
         is_auto_end (bool, optional): is this an automatic use of end_session and should be skipped with bypass_auto_end_session
     """
-    return Client().end_session(
+    Client().end_session(
         end_state=end_state,
         end_state_reason=end_state_reason,
         video=video,
diff --git a/agentops/enums.py b/agentops/enums.py
index 2c93be15..5a687862 100644
--- a/agentops/enums.py
+++ b/agentops/enums.py
@@ -9,20 +9,6 @@ class EventType(Enum):
     ERROR = "errors"
 
 
-class Models(Enum):
-    GPT_3_5_TURBO = "gpt-3.5-turbo"
-    GPT_3_5_TURBO_0301 = "gpt-3.5-turbo-0301"
-    GPT_3_5_TURBO_0613 = "gpt-3.5-turbo-0613"
-    GPT_3_5_TURBO_16K = "gpt-3.5-turbo-16k"
-    GPT_3_5_TURBO_16K_0613 = "gpt-3.5-turbo-16k-0613"
-    GPT_4_0314 = "gpt-4-0314"
-    GPT_4 = "gpt-4"
-    GPT_4_32K = "gpt-4-32k"
-    GPT_4_32K_0314 = "gpt-4-32k-0314"
-    GPT_4_0613 = "gpt-4-0613"
-    TEXT_EMBEDDING_ADA_002 = "text-embedding-ada-002"
-
-
 class EndState(Enum):
     SUCCESS = "Success"
     FAIL = "Fail"
diff --git a/agentops/event.py b/agentops/event.py
index 9bfe3a8b..a08efcd7 100644
--- a/agentops/event.py
+++ b/agentops/event.py
@@ -5,10 +5,10 @@
     Event: Represents discrete events to be recorded.
 """
 
-from dataclasses import asdict, dataclass, field
+from dataclasses import dataclass, field
 from typing import Any, Dict, List, Optional, Sequence, Union
 from .helpers import get_ISO_time, check_call_stack_for_agent_id
-from .enums import EventType, Models
+from .enums import EventType
 from uuid import UUID, uuid4
 import traceback
 
@@ -72,7 +72,7 @@ class LLMEvent(Event):
     prompt_tokens(int, optional): The number of tokens in the prompt message.
     completion(str, object, optional): The message or returned by the LLM. Preferably in ChatML format which is more fully supported by AgentOps.
     completion_tokens(int, optional): The number of tokens in the completion message.
-    model(Models, str, optional): LLM model e.g. "gpt-4". Models defined in enums.Models are more fully supported by AgentOps e.g. extra features in dashboard.
+    model(str, optional): LLM model e.g. "gpt-4", "gpt-3.5-turbo".
 
     """
 
@@ -82,7 +82,7 @@ class LLMEvent(Event):
     prompt_tokens: Optional[int] = None
     completion: Union[str, object] = None
     completion_tokens: Optional[int] = None
-    model: Optional[Union[Models, str]] = None
+    model: Optional[str] = None
 
 
 @dataclass