aiverify-moonshot 0.4.0__py3-none-any.whl

moonshot/integrations/web_api/services/benchmark_test_manager.py

@@ -0,0 +1,106 @@
+import asyncio
+import uuid
+from typing import Any
+
+from dependency_injector.wiring import inject
+
+from moonshot.src.runners.runner import Runner
+
+from ..schemas.benchmark_runner_dto import BenchmarkRunnerDTO
+from ..services.base_service import BaseService
+from ..services.benchmark_test_state import BenchmarkTestState
+from ..services.runner_service import RunnerService
+from ..status_updater.interface.benchmark_progress_callback import (
+    InterfaceBenchmarkProgressCallback,
+)
+from ..types.types import BenchmarkCollectionType, TestRunProgress
+
+
+class BenchmarkTestManager(BaseService):
+    @inject
+    def __init__(
+        self,
+        benchmark_test_state: BenchmarkTestState,
+        progress_status_updater: InterfaceBenchmarkProgressCallback,
+        runner_service: RunnerService,
+    ) -> None:
+        self.benchmark_test_state = benchmark_test_state
+        self.progress_status_updater = progress_status_updater
+        self.runner_service = runner_service
+        super().__init__()
+
+    def generate_unique_task_id(self) -> str:
+        unique_id = str(uuid.uuid4())
+        return f"task_{unique_id}"
+
+    def add_task(
+        self, executor_id: str, task: asyncio.Task[Any], moonshot_runner: Runner
+    ) -> None:
+        self.benchmark_test_state.add_task(executor_id, task, moonshot_runner)
+
+    def remove_task(self, executor_id: str) -> None:
+        self.benchmark_test_state.remove_task(executor_id)
+
+    async def cancel_task(self, executor_id: str) -> None:
+        await self.benchmark_test_state.cancel_task(executor_id)
+
+    def update_progress_status(self, updates: TestRunProgress):
+        self.benchmark_test_state.update_progress_status(updates)
+
+    def on_task_completed(self, task: asyncio.Task[Any]) -> None:
+        self.logger.debug(f"Task {task.get_name()} has completed")
+
+    async def run_test(
+        self,
+        benchmark_input_data: BenchmarkRunnerDTO,
+        benchmark_type: BenchmarkCollectionType,
+        moonshot_runner: Runner,
+    ) -> None:
+        try:
+            if benchmark_type == BenchmarkCollectionType.COOKBOOK:
+                async_run = moonshot_runner.run_cookbooks(
+                    cookbooks=benchmark_input_data.inputs,
+                    num_of_prompts=benchmark_input_data.num_of_prompts,
+                    random_seed=benchmark_input_data.random_seed,
+                    system_prompt=benchmark_input_data.system_prompt,
+                )
+            else:
+                async_run = moonshot_runner.run_recipes(
+                    recipes=benchmark_input_data.inputs,
+                    num_of_prompts=benchmark_input_data.num_of_prompts,
+                    random_seed=benchmark_input_data.random_seed,
+                    system_prompt=benchmark_input_data.system_prompt,
+                )
+        except Exception as e:
+            self.logger.error(f"Failed to execute benchmark - {e}")
+            raise Exception(f"Unexpected error in core library - {e}")
+
+        await async_run
+        async_run.close()
+
+    async def schedule_test_task(
+        self, input_data: BenchmarkRunnerDTO, benchmark_type: BenchmarkCollectionType
+    ) -> str:
+        task_id = self.generate_unique_task_id()
+        runner = self.runner_service.create_runner(
+            input_data.run_name,
+            input_data.endpoints,
+            input_data.description,
+            self.progress_status_updater.on_progress_update,
+        )
+        benchmark_coroutine = self.run_test(input_data, benchmark_type, runner)
+
+        task = asyncio.create_task(benchmark_coroutine, name=task_id)
+
+        def on_executor_completion(task: asyncio.Task[Any]):
+            if task.exception():
+                self.logger.error(
+                    f"Executor {task.get_name()} has failed - {task.exception()}"
+                )
+            else:
+                self.logger.debug(f"Executor {task.get_name()} has completed")
+
+        task.add_done_callback(on_executor_completion)
+
+        self.add_task(runner.id, task, runner)
+        return runner.id

moonshot/integrations/web_api/services/benchmark_test_state.py

@@ -0,0 +1,56 @@
+import asyncio
+from typing import Any, TypedDict
+
+from moonshot.integrations.web_api.services.base_service import BaseService
+from moonshot.integrations.web_api.types.types import TestRunProgress
+from moonshot.src.runners.runner import Runner
+
+
+class BenchmarkTaskInfo(TypedDict):
+    async_task: asyncio.Task[Any]
+    moonshot_runner: Runner
+    status: TestRunProgress | None
+
+
+class BenchmarkTestState(BaseService):
+    state: dict[str, BenchmarkTaskInfo] = {}
+
+    def add_task(
+        self, id: str, task: asyncio.Task[Any], moonshot_runner: Runner
+    ) -> None:
+        self.state[id] = {
+            "async_task": task,
+            "moonshot_runner": moonshot_runner,
+            "status": None,
+        }
+
+    def remove_task(self, id: str) -> BenchmarkTaskInfo:
+        return self.state.pop(id)
+
+    async def cancel_task(self, id: str) -> None:
+        task_info = self.state[id]
+        if task_info and not task_info["async_task"].done():
+            if task_info["moonshot_runner"]:
+                await task_info["moonshot_runner"].cancel()
+            task_info["async_task"].cancel()
+            self.logger.debug(f"Task {id} has been cancelled")
+        else:
+            self.logger.debug(f"Task {id} is already completed or does not exist.")
+
+    def update_progress_status(self, updates: TestRunProgress) -> None:
+        self.state[updates["current_runner_id"]]["status"] = updates
+
+    def get_state(self) -> dict[str, BenchmarkTaskInfo]:
+        return self.state
+
+    def get_progress_status(self, id: str) -> TestRunProgress | None:
+        if id in self.state:
+            return self.state[id]["status"]
+        return None
+
+    def get_all_progress_status(self) -> dict[str, TestRunProgress]:
+        return {
+            id: task_info["status"]
+            for id, task_info in self.state.items()
+            if task_info["status"] is not None
+        }

moonshot/integrations/web_api/services/benchmarking_service.py

@@ -0,0 +1,31 @@
+from dependency_injector.wiring import inject
+
+from ..schemas.benchmark_runner_dto import BenchmarkRunnerDTO
+from ..services.base_service import BaseService
+from ..services.benchmark_test_manager import BenchmarkTestManager
+from ..services.utils.exceptions_handler import exception_handler
+from ..types.types import BenchmarkCollectionType
+
+
+class BenchmarkingService(BaseService):
+    @inject
+    def __init__(self, benchmark_test_manager: BenchmarkTestManager) -> None:
+        self.benchmark_test_manager = benchmark_test_manager
+
+    @exception_handler
+    async def execute_cookbook(self, cookbook_executor_data: BenchmarkRunnerDTO) -> str:
+        id = await self.benchmark_test_manager.schedule_test_task(
+            cookbook_executor_data, BenchmarkCollectionType.COOKBOOK
+        )
+        return id
+
+    @exception_handler
+    async def execute_recipe(self, recipe_executor_data: BenchmarkRunnerDTO):
+        id = await self.benchmark_test_manager.schedule_test_task(
+            recipe_executor_data, BenchmarkCollectionType.RECIPE
+        )
+        return id
+
+    @exception_handler
+    async def cancel_executor(self, executor_id: str) -> None:
+        await self.benchmark_test_manager.cancel_task(executor_id)

moonshot/integrations/web_api/services/context_strategy_service.py

@@ -0,0 +1,22 @@
+from typing import Any
+
+from .... import api as moonshot_api
+from ..services.utils.exceptions_handler import exception_handler
+from .base_service import BaseService
+
+
+class ContextStrategyService(BaseService):
+
+    @exception_handler
+    def get_ctx_strategy(self) -> list[dict]:
+        strategies = moonshot_api.api_get_all_context_strategy_metadata()
+        return strategies
+    
+    @exception_handler
+    def get_ctx_strategy_name(self) -> list[str]:
+        strategies = moonshot_api.api_get_all_context_strategies()
+        return strategies
+
+    @exception_handler
+    def delete_ctx_strategy(self, ctx_strategy_name: str) -> None:
+        moonshot_api.api_delete_context_strategy(ctx_strategy_name)

moonshot/integrations/web_api/services/cookbook_service.py

@@ -0,0 +1,194 @@
+from moonshot.src.cookbooks.cookbook_arguments import CookbookArguments as Cookbook
+from moonshot.src.recipes.recipe_arguments import RecipeArguments as Recipe
+
+from .... import api as moonshot_api
+from ..schemas.cookbook_create_dto import CookbookCreateDTO, CookbookUpdateDTO
+from ..schemas.cookbook_response_model import CookbookResponseModel
+from ..services.base_service import BaseService
+from ..services.recipe_service import get_total_prompt_in_recipe
+from ..services.utils.exceptions_handler import exception_handler
+
+
+class CookbookService(BaseService):
+    @exception_handler
+    def create_cookbook(self, cookbook_data: CookbookCreateDTO) -> None:
+        """
+        Create a new cookbook with the given data.
+
+        Args:
+            cookbook_data (CookbookCreateDTO): Data transfer object containing cookbook details.
+        """
+        moonshot_api.api_create_cookbook(
+            name=cookbook_data.name,
+            description=cookbook_data.description,
+            recipes=cookbook_data.recipes,
+        )
+
+    @exception_handler
+    def get_all_cookbooks(
+        self,
+        tags: str,
+        categories: str,
+        count: bool,
+        ids: str | None = None,
+        categories_excluded: str | None = None,
+    ) -> list[CookbookResponseModel]:
+        """
+        Retrieve all cookbooks, optionally filtered by tags and categories, and with prompt counts.
+
+        Args:
+            ids (str): Filter cookbooks by IDs.
+            tags (str): Filter cookbooks by tags.
+            categories (str): Filter cookbooks by categories.
+            count (bool): Include the total prompt count in each cookbook.
+
+        Returns:
+            list[CookbookResponseModel]: A list of cookbook response models.
+        """
+        cookbooks_list: list[CookbookResponseModel] = []
+
+        if ids:
+            cookbook_ids = ids.split(",")
+            cookbooks = [
+                moonshot_api.api_read_cookbook(cookbook_id)
+                for cookbook_id in cookbook_ids
+            ]
+        else:
+            cookbooks = moonshot_api.api_get_all_cookbook()
+
+        for cookbook_dict in cookbooks:
+            cookbook = CookbookResponseModel(**cookbook_dict)
+
+            if not tags and not categories:
+                if cookbook not in cookbooks_list:
+                    cookbooks_list.append(cookbook)
+                    if count:
+                        cookbook.total_prompt_in_cookbook = (
+                            get_total_prompt_in_cookbook(cookbook)
+                        )
+
+            if tags and cookbooks_recipe_has_tags(tags, cookbook):
+                if cookbook not in cookbooks_list:
+                    cookbooks_list.append(cookbook)
+                    if count:
+                        cookbook.total_prompt_in_cookbook = (
+                            get_total_prompt_in_cookbook(cookbook)
+                        )
+
+            if categories and cookbooks_recipe_has_categories(categories, cookbook):
+                if cookbook not in cookbooks_list:
+                    cookbooks_list.append(cookbook)
+                    if count:
+                        cookbook.total_prompt_in_cookbook = (
+                            get_total_prompt_in_cookbook(cookbook)
+                        )
+
+            if categories_excluded and cookbooks_recipe_has_categories(
+                categories_excluded, cookbook
+            ):
+                cookbooks_list.remove(cookbook)
+
+        return cookbooks_list
+
+    @exception_handler
+    def get_all_cookbooks_names(self) -> list[str]:
+        """
+        Retrieve the names of all cookbooks.
+
+        Returns:
+            list[str]: A list of cookbook names.
+        """
+        cookbooks = moonshot_api.api_get_all_cookbook_name()
+        return cookbooks
+
+    @exception_handler
+    def update_cookbook(
+        self, cookbook_data: CookbookUpdateDTO, cookbook_id: str
+    ) -> None:
+        """
+        Update an existing cookbook with new data.
+
+        Args:
+            cookbook_data (CookbookCreateDTO): Data transfer object containing new cookbook details.
+            cookbook_id (str): The ID of the cookbook to update.
+        """
+        update_data = {
+            k: v
+            for k, v in cookbook_data.to_dict().items()
+            if v is not None and k != "id"
+        }
+        moonshot_api.api_update_cookbook(cb_id=cookbook_id, **update_data)
+
+    @exception_handler
+    def delete_cookbook(self, cookbook_id: str) -> None:
+        """
+        Delete a cookbook by its ID.
+
+        Args:
+            cookbook_id (str): The ID of the cookbook to delete.
+        """
+        moonshot_api.api_delete_cookbook(cookbook_id)
+
+
+@staticmethod
+def get_total_prompt_in_cookbook(cookbook: Cookbook) -> int:
+    """
+    Calculate the total number of prompts in a cookbook.
+
+    This function sums up the total prompts for each recipe in the cookbook.
+
+    Args:
+        cookbook (Cookbook): The cookbook object containing the recipe IDs.
+
+    Returns:
+        int: The total count of prompts within the cookbook.
+    """
+    recipes = moonshot_api.api_read_recipes(cookbook.recipes)
+    return sum(get_total_prompt_in_recipe(Recipe(**recipe)) for recipe in recipes)
+
+
+@staticmethod
+def cookbooks_recipe_has_tags(tags: str, cookbook: Cookbook) -> bool:
+    """
+    Check if any recipe in a cookbook has the specified tags.
+
+    Args:
+        tags (str): The tags to check for in the cookbook's recipes.
+        cookbook (Cookbook): The cookbook object containing the recipe IDs.
+
+    Returns:
+        bool: True if any recipe in the cookbook has the specified tags, False otherwise.
+    """
+    recipe_ids = cookbook.recipes
+    recipes = moonshot_api.api_read_recipes(recipe_ids)
+    for recipe in recipes:
+        recipe = Recipe(**recipe)
+        if tags in recipe.tags:
+            return True
+    return False
+
+
+@staticmethod
+def cookbooks_recipe_has_categories(categories: str, cookbook: Cookbook) -> bool:
+    """
+    Check if any recipe in a cookbook has the specified categories.
+
+    Args:
+        categories (str): The categories to check for in the cookbook's recipes.
+        cookbook (Cookbook): The cookbook object containing the recipe IDs.
+        exclude_categories (str): The categories to exclude
+
+    Returns:
+        bool: True if any recipe in the cookbook has the specified categories, False otherwise.
+    """
+    recipe_ids = cookbook.recipes
+    categories_list = [category.lower() for category in categories.split(",")]
+    recipes = moonshot_api.api_read_recipes(recipe_ids)
+    for recipe in recipes:
+        recipe = Recipe(**recipe)
+        if any(
+            category in [rcat.lower() for rcat in recipe.categories]
+            for category in categories_list
+        ):
+            return True
+    return False

moonshot/integrations/web_api/services/dataset_service.py

@@ -0,0 +1,20 @@
+from .... import api as moonshot_api
+from ..schemas.dataset_response_dto import DatasetResponseDTO
+from ..services.base_service import BaseService
+from ..services.utils.exceptions_handler import exception_handler
+
+
+class DatasetService(BaseService):
+    @exception_handler
+    def get_all_datasets(self) -> list[DatasetResponseDTO]:
+        datasets = moonshot_api.api_get_all_datasets()
+        return [DatasetResponseDTO.model_validate(dataset) for dataset in datasets]
+
+    @exception_handler
+    def get_all_datasets_name(self) -> list[str] | None:
+        datasets = moonshot_api.api_get_all_datasets_name()
+        return datasets
+
+    @exception_handler
+    def delete_dataset(self, dataset_id: str) -> None:
+        moonshot_api.api_delete_dataset(dataset_id)

moonshot/integrations/web_api/services/endpoint_service.py

@@ -0,0 +1,65 @@
+from .... import api as moonshot_api
+from ..schemas.endpoint_create_dto import EndpointUpdateDTO
+from ..schemas.endpoint_response_model import EndpointDataModel
+from .base_service import BaseService
+from .utils.exceptions_handler import exception_handler
+
+
+class EndpointService(BaseService):
+    @exception_handler
+    def add_endpoint(self, endpoint_data: EndpointDataModel) -> None:
+        moonshot_api.api_create_endpoint(
+            name=endpoint_data.name,
+            connector_type=endpoint_data.connector_type,
+            uri=endpoint_data.uri,
+            token=endpoint_data.token,
+            max_calls_per_second=endpoint_data.max_calls_per_second,
+            max_concurrency=endpoint_data.max_concurrency,
+            params=endpoint_data.params,
+        )
+
+    @exception_handler
+    def get_all_endpoints(self) -> list[EndpointDataModel | None]:
+        filtered_endpoints = []
+        endpoints = moonshot_api.api_get_all_endpoint()
+        for endpoint in endpoints:
+            filtered_endpoints.append(EndpointDataModel.model_validate(endpoint))
+
+        filtered_endpoints.sort(
+            key=lambda endpoint: endpoint.created_date, reverse=True
+        )
+
+        return [
+            EndpointDataModel.model_validate(endpoint)
+            for endpoint in filtered_endpoints
+        ]
+
+    @exception_handler
+    def get_endpoint(self, endpoint_id: str) -> EndpointDataModel | None:
+        endpoint = moonshot_api.api_read_endpoint(endpoint_id)
+        return EndpointDataModel.model_validate(endpoint)
+
+    @exception_handler
+    def get_all_endpoints_names(self) -> list[str]:
+        endpoints_names = moonshot_api.api_get_all_endpoint_name()
+        return endpoints_names
+
+    @exception_handler
+    def update_endpoint(
+        self, endpoint_id: str, endpoint_data: EndpointUpdateDTO
+    ) -> None:
+        update_data = {
+            k: v
+            for k, v in endpoint_data.to_dict().items()
+            if v is not None and k != "id"
+        }
+        moonshot_api.api_update_endpoint(ep_id=endpoint_id, **update_data)
+
+    @exception_handler
+    def delete_endpoint(self, endpoint_id: str) -> None:
+        moonshot_api.api_delete_endpoint(endpoint_id)
+
+    @exception_handler
+    def get_all_connectors(self) -> list[str]:
+        connectors = moonshot_api.api_get_all_connector_type()
+        return connectors

moonshot/integrations/web_api/services/metric_service.py

@@ -0,0 +1,14 @@
+from .... import api as moonshot_api
+from ..services.base_service import BaseService
+from ..services.utils.exceptions_handler import exception_handler
+
+
+class MetricService(BaseService):
+    @exception_handler
+    def get_all_metric(self) -> list[dict]:
+        metrics = moonshot_api.api_get_all_metric()
+        return metrics
+
+    @exception_handler
+    def delete_metric(self, metric_id: str) -> None:
+        moonshot_api.api_delete_metric(metric_id)

moonshot/integrations/web_api/services/prompt_template_service.py

@@ -0,0 +1,39 @@
+from typing import Any
+
+from .... import api as moonshot_api
+from ..services.utils.exceptions_handler import exception_handler
+from .base_service import BaseService
+
+
+class PromptTemplateService(BaseService):
+    @exception_handler
+    def get_prompt_templates(self) -> list[dict[str, Any]]:
+        """
+        Retrieve all prompt templates with their details.
+
+        Returns:
+            list[dict[str, Any]]: A list of dictionaries containing the details of each prompt template.
+        """
+        templates = moonshot_api.api_get_all_prompt_template_detail()
+        return templates
+
+    @exception_handler
+    def get_prompt_templates_name(self) -> list[str]:
+        """
+        Retrieve the names of all prompt templates.
+
+        Returns:
+            list[str]: A list of names of all prompt templates.
+        """
+        templates = moonshot_api.api_get_all_prompt_template_name()
+        return templates
+
+    @exception_handler
+    def delete_prompt_template(self, pt_id: str):
+        """
+        Delete a prompt template by its identifier.
+
+        Args:
+            pt_id (str): The unique identifier of the prompt template to be deleted.
+        """
+        moonshot_api.api_delete_prompt_template(pt_id)