Adding an ENUM for application stage (#3)

* Adding an ENUM for application stage

* Fixed bug in dataset creation and updated the notebook with the Enum utility

---------

Co-authored-by: Preetam Joshi <info@aimon.ai>

Author: pjoshi30

aimon/__init__.py

@@ -1,4 +1,5 @@
 from .client import Client
+from .client import ApplicationStage
 from .dataset import Dataset, DatasetCollection, DatasetRecord
 from .evaluation import Evaluation, Run
 from .models import MLModel

aimon/client.py

@@ -10,6 +10,10 @@
 class InvalidAPIKeyError(Exception):
     pass
 
+# An Enum for the different stages of an application
+class ApplicationStage:
+    PRODUCTION = "production"
+    EVALUATION = "evaluation"
 
 class Client(object):
 
@@ -80,11 +84,13 @@ def application(self, name, model, stage, app_type, metadata=None):
         Gets or creates an Application object.
         :param name: The name of the model
         :param model: An MLModel object
-        :param stage: The stage of the application. Should be either "production" or "experimentation"
+        :param stage: The stage of the application. Should be either "production" or "evaluation"
         :param app_type: The type of the application
         :param metadata: The metadata associated with the model
         """
         headers = self.create_header()
+        if stage.lower() not in ['production', 'evaluation']:
+            raise Exception("Invalid stage specified. Should be either 'production' or 'evaluation'")
         data = {
             "name": name,
             "model_name": model.name,
@@ -173,8 +179,8 @@ def create_dataset(self, name, file_path, description):
         dataset = get_request('{}/v1/dataset'.format(AIMON_SDK_BACKEND_URL), headers=headers,
                               params={'name': name})
         if dataset:
-            return Dataset(dataset['name'], dataset['description'], dataset['creation_time'],
-                           dataset['last_updated_time'], dataset['s3_location'], dataset['sha'], dataset['user_id'])
+            return Dataset(self.api_key, dataset['name'], dataset['description'], dataset['creation_time'],
+                           dataset['last_updated_time'], dataset['sha'], dataset['user_id'])
 
         # Dataset does not exist, create a new one
         if not file_path:

examples/notebooks/aimon_sdk_langchain_summarization.ipynb

@@ -29,7 +29,7 @@
    "source": [
     "# Create the AIMon client. You would need an API Key (that can be retrieved from the UI in your user profile). \n",
     "# Your user email is also needed to track ownership of components such as datasets and applications.\n",
-    "aimon_client = Client(api_key=\"YOUR API KEY HERE\", email=\"YOUR EMAIL ADDRESS HERE\")"
+    "aimon_client = Client(api_key=\"YOUR AIMon API KEY HERE\", email=\"YOUR EMAIL ID HERE\")"
    ]
   },
   {
@@ -52,7 +52,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "['GPT-4']\n"
+      "['GPT-4', 'Llama-2', 'model_type']\n"
      ]
     }
    ],
@@ -92,16 +92,17 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 5,
    "id": "22794a7c-8193-4c10-b633-fe4d0d4277d0",
    "metadata": {},
    "outputs": [],
    "source": [
+    "from aimon import ApplicationStage\n",
     "# Using the AIMon client, create or get an existing application\n",
     "new_app = aimon_client.application(\n",
     "    \"my_llm_summarization_app\", \n",
     "    my_model, \n",
-    "    stage=\"experimentation\", \n",
+    "    stage=ApplicationStage.EVALUATION, \n",
     "    app_type=\"summarization\", \n",
     "    metadata={\"applicaiton_url\": \"https://acme.com/summarization\"}\n",
     ")"
@@ -119,7 +120,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 6,
    "id": "af98abfc-c04d-4cc7-ba3c-62f550de0c99",
    "metadata": {},
    "outputs": [],
@@ -177,19 +178,19 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 7,
    "id": "34881793-3778-4114-8066-7c41cfb6d462",
    "metadata": {},
    "outputs": [],
    "source": [
     "# Create a new dataset\n",
-    "# dataset1 = aimon_client.create_dataset(\"test_evaluation_dataset_1.csv\", \"./aimon_sdk_files/test_evaluation_dataset_1.csv\", \"This is one custom dataset\")\n",
-    "# dataset2 = aimon_client.create_dataset(\"test_evaluation_dataset_2.csv\", \"./aimon_sdk_files/test_evaluation_dataset_2.csv\", \"This is another custom dataset\")"
+    "dataset1 = aimon_client.create_dataset(\"test_evaluation_dataset_1.csv\", \"/Users/preetamjoshi/Downloads/test_evaluation_dataset_1.csv\", \"This is one custom dataset\")\n",
+    "dataset2 = aimon_client.create_dataset(\"test_evaluation_dataset_2.csv\", \"/Users/preetamjoshi/Downloads/test_evaluation_dataset_2.csv\", \"This is another custom dataset\")"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 8,
    "id": "c7cc2795-6010-479a-90c0-bd275589f0cf",
    "metadata": {},
    "outputs": [],
@@ -210,7 +211,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 9,
    "id": "0dbceabd-a6a2-4671-959d-ca15025841b1",
    "metadata": {},
    "outputs": [],
@@ -235,7 +236,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 10,
    "id": "0f7991bf-a7b5-42f6-9a45-08c66c98d41c",
    "metadata": {},
    "outputs": [],
@@ -259,7 +260,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 11,
    "id": "7b76a493-8738-40ba-a722-e22425c37e1f",
    "metadata": {},
    "outputs": [],
@@ -274,17 +275,17 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 12,
    "id": "c2706102-672d-4e11-b77b-c1110a31772f",
    "metadata": {},
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "/Users/preetamjoshi/projects/aimon/aimon/datascience/aimon_dev/lib/python3.10/site-packages/langchain_core/_api/deprecation.py:117: LangChainDeprecationWarning: The class `langchain_community.llms.openai.OpenAI` was deprecated in langchain-community 0.0.10 and will be removed in 0.2.0. An updated version of the class exists in the langchain-openai package and should be used instead. To use it run `pip install -U langchain-openai` and import as `from langchain_openai import OpenAI`.\n",
+      "/Users/preetamjoshi/projects/aimon/aimon-rely/artry/lib/python3.9/site-packages/langchain_core/_api/deprecation.py:119: LangChainDeprecationWarning: The class `OpenAI` was deprecated in LangChain 0.0.10 and will be removed in 0.3.0. An updated version of the class exists in the langchain-openai package and should be used instead. To use it run `pip install -U langchain-openai` and import as `from langchain_openai import OpenAI`.\n",
       "  warn_deprecated(\n",
-      "/Users/preetamjoshi/projects/aimon/aimon/datascience/aimon_dev/lib/python3.10/site-packages/langchain_core/_api/deprecation.py:117: LangChainDeprecationWarning: The function `run` was deprecated in LangChain 0.1.0 and will be removed in 0.2.0. Use invoke instead.\n",
+      "/Users/preetamjoshi/projects/aimon/aimon-rely/artry/lib/python3.9/site-packages/langchain_core/_api/deprecation.py:119: LangChainDeprecationWarning: The method `Chain.run` was deprecated in langchain 0.1.0 and will be removed in 0.3.0. Use invoke instead.\n",
       "  warn_deprecated(\n"
      ]
     },
@@ -332,7 +333,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 15,
    "id": "bb81a8ad-e10a-43f5-9682-844d9ab2ccb4",
    "metadata": {},
    "outputs": [],
@@ -352,7 +353,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 16,
    "id": "1e0dc9d7-e0b0-4a05-8967-0e06179486a3",
    "metadata": {},
    "outputs": [
@@ -394,7 +395,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.13"
+   "version": "3.9.0"
   }
  },
  "nbformat": 4,