Enhancement - add prefilled testsets for apps created with template i…

…n demo version (#623)

add pre-filled testsets to apps created from templates

agenta-backend/agenta_backend/resources/default_testsets/single_prompt_testsets.json

@@ -0,0 +1,9 @@
+[
+	{"country": "Nauru", "correct_answer": "The capital of Nauru is Funafuti"},
+	{"country": "Tuvalu", "correct_answer": "The capital of Tuvalu is Funafuti"},
+	{"country": "Brunei", "correct_answer": "The capital of Brunei is Bandar Seri Begawan"},
+	{"country": "Kiribati", "correct_answer": "The capital of Kiribati is Tarawa"},
+	{"country": "Comoros", "correct_answer": "The capital of Comoros is Moroni"},
+	{"country": "Kyrgyzstan", "correct_answer": "The capital of Kyrgyzstan is Bishkek"},
+	{"country": "Azerbaijan", "correct_answer": "The capital of Azerbaijan is Baku"}
+]

agenta-backend/agenta_backend/routers/app_variant.py

@@ -264,7 +264,9 @@ async def stop_variant(app_variant: AppVariant):
 
 
 @router.get("/list_images/", response_model=List[Image])
-async def list_images(stoken_session: SessionContainer = Depends(verify_session())):
+async def list_images(
+    stoken_session: SessionContainer = Depends(verify_session()),
+):
     """Lists the images from our repository
 
     Raises:
@@ -436,12 +438,15 @@ async def add_app_variant_from_template(
     if variant_exist is None:
         # Save variant based on the image to database
         await db_manager.add_variant_based_on_image(app_variant, image, **kwargs)
+
+        # Create testset for apps created
+        await db_manager.add_testset_to_app_variant(app_variant, image, **kwargs)
     else:
         # Update variant based on the image
         await app_manager.update_variant_image(app_variant, image, **kwargs)
 
     # Start variant
-    url = await app_manager.start_variant(app_variant, envvars, **kwargs)
+    await app_manager.start_variant(app_variant, envvars, **kwargs)
 
     return {
         "message": "Variant created and running!",

agenta-backend/agenta_backend/services/db_manager.py

@@ -1,6 +1,9 @@
-import logging
 import os
-from typing import Any, Dict, List
+import logging
+from pathlib import Path
+from bson import ObjectId
+from datetime import datetime
+from typing import Dict, List, Any
 
 from agenta_backend.models.api.api_models import (
     App,
@@ -15,13 +18,15 @@
     image_db_to_pydantic,
     templates_db_to_pydantic,
 )
+from agenta_backend.services.json_importer_helper import get_json
 from agenta_backend.models.db_engine import DBEngine
 from agenta_backend.models.db_models import (
     AppVariantDB,
     EnvironmentDB,
     ImageDB,
     OrganizationDB,
     TemplateDB,
+    TestSetDB,
     UserDB,
 )
 from agenta_backend.services import helpers
@@ -34,6 +39,9 @@
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
 
+# Define parent directory
+PARENT_DIRECTORY = Path(os.path.dirname(__file__)).parent
+
 
 async def get_templates() -> List[Template]:
     templates = await engine.find(TemplateDB)
@@ -117,6 +125,34 @@ async def add_variant_based_on_image(
     await engine.save(db_app_variant)
 
 
+async def add_testset_to_app_variant(
+    app_variant: AppVariant, image: Image, **kwargs: dict
+):
+    """Add testset to app variant.
+
+    Args:
+        app_variant (AppVariant): the app variant
+        image (Image): the image
+    """
+
+    user_instance = await get_user_object(kwargs["uid"])
+
+    app_template_name = image.tags.split(":")[-1]
+    if app_template_name == "single_prompt":
+        json_path = (
+            f"{PARENT_DIRECTORY}/resources/default_testsets/single_prompt_testsets.json"
+        )
+        csvdata = get_json(json_path)
+        testset = {
+            "name": f"{app_variant.app_name}_testset",
+            "app_name": app_variant.app_name,
+            "created_at": datetime.now().isoformat(),
+            "csvdata": csvdata,
+        }
+        testset = TestSetDB(**testset, user=user_instance)
+        await engine.save(testset)
+
+
 async def add_variant_based_on_previous(
     previous_app_variant: AppVariant,
     new_variant_name: str,
@@ -237,7 +273,10 @@ async def list_app_variants(
 
 
 async def get_app_variant_by_app_name_and_variant_name(
-    app_name: str, variant_name: str, show_soft_deleted: bool = False, **kwargs: dict
+    app_name: str,
+    variant_name: str,
+    show_soft_deleted: bool = False,
+    **kwargs: dict,
 ) -> AppVariant:
     """Fetches an app variant based on app_name and variant_name.
 
@@ -489,7 +528,9 @@ async def remove_image(image: ImageExtended, **kwargs: dict):
     await engine.delete(image_db)
 
 
-async def check_is_last_variant_for_image(db_app_variant: AppVariantDB) -> bool:
+async def check_is_last_variant_for_image(
+    db_app_variant: AppVariantDB,
+) -> bool:
     """Checks whether the input variant is the sole variant that uses its linked image
     This is a helpful function to determine whether to delete the image when removing a variant
     Usually many variants will use the same image (these variants would have been created using the UI)

agenta-backend/agenta_backend/services/json_importer_helper.py

@@ -0,0 +1,17 @@
+import json
+
+
+def get_json(json_path: str):
+    """Reads and returns the contents of a JSON file as a list of
+    dictionaries.
+
+    Args:
+                json_path (str): The path of json
+    """
+
+    with open(json_path) as f:
+        try:
+            json_data = json.loads(f.read())
+        except Exception:
+            raise ValueError(f"Could not read JSON file: {json_path}")
+    return json_data