dataset migration

athina-ai · May 16, 2024 · 57ed947 · 57ed947
1 parent b6aa8f0
commit 57ed947
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 6 deletions.
diff --git a/athina/datasets/dataset.py b/athina/datasets/dataset.py
@@ -15,6 +15,7 @@ class DatasetRow:
 class Dataset:
     id: str
     source: str
+    data_source: str
     name: str
     description: Optional[str] = None
     language_model_id: Optional[str] = None
@@ -39,6 +40,7 @@ def create(name: str, description: Optional[str] = None, language_model_id: Opti
         """
         dataset_data = {
                 "source": "dev_sdk",
+                "data_source": "dev_sdk",
                 "name": name,
                 "description": description,
                 "language_model_id": language_model_id,
@@ -53,7 +55,7 @@ def create(name: str, description: Optional[str] = None, language_model_id: Opti
             created_dataset_data = AthinaApiService.create_dataset(dataset_data)
         except Exception as e:
             raise
-        dataset = Dataset(id=created_dataset_data['id'], source=created_dataset_data['source'], name=created_dataset_data['name'], description=created_dataset_data['description'], language_model_id=created_dataset_data['language_model_id'], prompt_template=created_dataset_data['prompt_template'])
+        dataset = Dataset(id=created_dataset_data['id'], source=created_dataset_data['source'], data_source=created_dataset_data['data_source'], name=created_dataset_data['name'], description=created_dataset_data['description'], language_model_id=created_dataset_data['language_model_id'], prompt_template=created_dataset_data['prompt_template'])
         return dataset
 
     @staticmethod

diff --git a/athina/services/athina_api_service.py b/athina/services/athina_api_service.py
@@ -129,7 +129,7 @@ def create_dataset(
         Creates a dataset by calling the Athina API
         """
         try:
-            endpoint = f"{API_BASE_URL}/api/v1/dataset"
+            endpoint = f"{API_BASE_URL}/api/v1/dataset_v2"
             response = requests.post(
                 endpoint,
                 headers=AthinaApiService._headers(),
@@ -166,7 +166,7 @@ def add_dataset_rows(dataset_id: str, rows: List[Dict]):
         - CustomException: If the API call fails or returns an error.
         """
         try:
-            endpoint = f"{API_BASE_URL}/api/v1/dataset/{dataset_id}/add-rows"
+            endpoint = f"{API_BASE_URL}/api/v1/dataset_v2/{dataset_id}/add-rows"
             response = requests.post(
                 endpoint,
                 headers=AthinaApiService._headers(),

diff --git a/examples/dataset_creation.ipynb b/examples/dataset_creation.ipynb
@@ -2,9 +2,20 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/Users/akshat_g/athina/repos/athina-evals/.venv/lib/python3.9/site-packages/urllib3/__init__.py:35: NotOpenSSLWarning: urllib3 v2 only supports OpenSSL 1.1.1+, currently the 'ssl' module is compiled with 'LibreSSL 2.8.3'. See: https://github.com/urllib3/urllib3/issues/3020\n",
+      "  warnings.warn(\n",
+      "/Users/akshat_g/athina/repos/athina-evals/.venv/lib/python3.9/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n"
+     ]
+    }
+   ],
    "source": [
     "import os\n",
     "from athina.datasets import Dataset\n",
@@ -17,7 +28,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -39,6 +50,15 @@
     "    print(f\"Failed to create dataset: {e}\")"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(dataset.id)   "
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,