Merge to test

aixplain/cli_groups.py

@@ -21,28 +21,43 @@
     CLI Runner
 """
 import click
-from aixplain.factories.cli.model_factory_cli import list_host_machines, list_functions, create_asset_repo, asset_repo_login, onboard_model, deploy_huggingface_model, get_huggingface_model_status, list_gpus
+from aixplain.factories.cli.model_factory_cli import (
+    list_host_machines,
+    list_functions,
+    create_asset_repo,
+    asset_repo_login,
+    onboard_model,
+    deploy_huggingface_model,
+    get_huggingface_model_status,
+    list_gpus,
+)
 
-@click.group('cli')
+
+@click.group("cli")
 def cli():
     pass
 
-@click.group('list')
+
+@click.group("list")
 def list():
     pass
 
-@click.group('get')
+
+@click.group("get")
 def get():
     pass
 
-@click.group('create')
+
+@click.group("create")
 def create():
     pass
 
-@click.group('onboard')
+
+@click.group("onboard")
 def onboard():
     pass
 
+
 cli.add_command(list)
 cli.add_command(get)
 cli.add_command(create)

aixplain/enums/__init__.py

@@ -13,3 +13,4 @@
 from .supplier import Supplier
 from .sort_by import SortBy
 from .sort_order import SortOrder
+from .model_status import ModelStatus

aixplain/enums/asset_status.py

@@ -24,20 +24,22 @@
 from enum import Enum
 from typing import Text
 
+
 class AssetStatus(Text, Enum):
-    HIDDEN = 'hidden'
-    SCHEDULED = 'scheduled'
-    ONBOARDING = 'onboarding'
-    ONBOARDED = 'onboarded'
-    PENDING = 'pending'
-    FAILED = 'failed'
-    TRAINING = 'training'
-    REJECTED = 'rejected'
-    ENABLING = 'enabling'
-    DELETING = 'deleting'
-    DISABLED = 'disabled'
-    DELETED = 'deleted'
-    IN_PROGRESS = 'in_progress'
-    COMPLETED = 'completed'
-    CANCELING = 'canceling'
-    CANCELED = 'canceled'
+    DRAFT = "draft"
+    HIDDEN = "hidden"
+    SCHEDULED = "scheduled"
+    ONBOARDING = "onboarding"
+    ONBOARDED = "onboarded"
+    PENDING = "pending"
+    FAILED = "failed"
+    TRAINING = "training"
+    REJECTED = "rejected"
+    ENABLING = "enabling"
+    DELETING = "deleting"
+    DISABLED = "disabled"
+    DELETED = "deleted"
+    IN_PROGRESS = "in_progress"
+    COMPLETED = "completed"
+    CANCELING = "canceling"
+    CANCELED = "canceled"

aixplain/enums/data_subtype.py

@@ -32,6 +32,6 @@ class DataSubtype(Enum):
     RACE = "race"
     SPLIT = "split"
     TOPIC = "topic"
-    
+
     def __str__(self):
         return self._value_

aixplain/enums/function.py

@@ -55,7 +55,7 @@ def load_functions():
                 if input_data_object["required"] is True
             },
             "output": {output_data_object["dataType"] for output_data_object in function["output"]},
-            "spec": function
+            "spec": function,
         }
         for function in resp["items"]
     }

aixplain/enums/model_status.py

@@ -0,0 +1,11 @@
+from enum import Enum
+from typing import Text
+
+
+class ModelStatus(Text, Enum):
+    FAILED = "FAILED"
+    IN_PROGRESS = "IN_PROGRESS"
+    SUCCESS = "SUCCESS"
+
+    def __str__(self):
+        return self._value_

aixplain/enums/storage_type.py

@@ -30,4 +30,4 @@ class StorageType(Enum):
     FILE = "file"
 
     def __str__(self):
-        return self._value_
+        return self._value_

aixplain/factories/agent_factory/__init__.py

@@ -34,7 +34,6 @@
 from aixplain.utils import config
 from typing import Dict, List, Optional, Text, Union
 
-from aixplain.factories.agent_factory.utils import build_agent, validate_llm, validate_name
 from aixplain.utils.file_utils import _request_with_retry
 from urllib.parse import urljoin
 
@@ -65,74 +64,49 @@ def create(
         Returns:
             Agent: created Agent
         """
-        validate_name(name)
-        # validate LLM ID
-        validate_llm(llm_id)
+        from aixplain.factories.agent_factory.utils import build_agent
 
+        agent = None
+        url = urljoin(config.BACKEND_URL, "sdk/agents")
+        headers = {"x-api-key": api_key}
+
+        if isinstance(supplier, dict):
+            supplier = supplier["code"]
+        elif isinstance(supplier, Supplier):
+            supplier = supplier.value["code"]
+
+        payload = {
+            "name": name,
+            "assets": [tool.to_dict() for tool in tools],
+            "description": description,
+            "supplier": supplier,
+            "version": version,
+            "llmId": llm_id,
+            "status": "draft",
+        }
+        agent = build_agent(payload=payload, api_key=api_key)
+        agent.validate()
+        response = "Unspecified error"
         try:
-            agent = None
-            url = urljoin(config.BACKEND_URL, "sdk/agents")
-            headers = {"x-api-key": api_key}
-
-            if isinstance(supplier, dict):
-                supplier = supplier["code"]
-            elif isinstance(supplier, Supplier):
-                supplier = supplier.value["code"]
-
-            tool_payload = []
-            for tool in tools:
-                if isinstance(tool, ModelTool):
-                    tool.validate()
-                    tool_payload.append(
-                        {
-                            "function": tool.function.value if tool.function is not None else None,
-                            "type": "model",
-                            "description": tool.description,
-                            "supplier": tool.supplier.value["code"] if tool.supplier else None,
-                            "version": tool.version if tool.version else None,
-                            "assetId": tool.model,
-                        }
-                    )
-                elif isinstance(tool, PipelineTool):
-                    tool.validate()
-                    tool_payload.append(
-                        {
-                            "assetId": tool.pipeline,
-                            "description": tool.description,
-                            "type": "pipeline",
-                        }
-                    )
-                else:
-                    raise Exception("Agent Creation Error: Tool type not supported.")
-
-            payload = {
-                "name": name,
-                "assets": tool_payload,
-                "description": description,
-                "supplier": supplier,
-                "version": version,
-                "llmId": llm_id,
-            }
-
-            logging.info(f"Start service for POST Create Agent  - {url} - {headers} - {json.dumps(payload)}")
+            logging.debug(f"Start service for POST Create Agent  - {url} - {headers} - {json.dumps(payload)}")
             r = _request_with_retry("post", url, headers=headers, json=payload)
-            if 200 <= r.status_code < 300:
-                response = r.json()
-                agent = build_agent(payload=response, api_key=api_key)
-            else:
-                error = r.json()
-                error_msg = "Agent Onboarding Error: Please contact the administrators."
-                if "message" in error:
-                    msg = error["message"]
-                    if error["message"] == "err.name_already_exists":
-                        msg = "Agent name already exists."
-                    elif error["message"] == "err.asset_is_not_available":
-                        msg = "Some tools are not available."
-                    error_msg = f"Agent Onboarding Error (HTTP {r.status_code}): {msg}"
-                logging.exception(error_msg)
-                raise Exception(error_msg)
-        except Exception as e:
-            raise Exception(e)
+            response = r.json()
+        except Exception:
+            raise Exception("Agent Onboarding Error: Please contact the administrators.")
+
+        if 200 <= r.status_code < 300:
+            agent = build_agent(payload=response, api_key=api_key)
+        else:
+            error_msg = f"Agent Onboarding Error: {response}"
+            if "message" in response:
+                msg = response["message"]
+                if response["message"] == "err.name_already_exists":
+                    msg = "Agent name already exists."
+                elif response["message"] == "err.asset_is_not_available":
+                    msg = "Some tools are not available."
+                error_msg = f"Agent Onboarding Error (HTTP {r.status_code}): {msg}"
+            logging.exception(error_msg)
+            raise Exception(error_msg)
         return agent
 
     @classmethod
@@ -141,6 +115,7 @@ def create_model_tool(
         model: Optional[Union[Model, Text]] = None,
         function: Optional[Union[Function, Text]] = None,
         supplier: Optional[Union[Supplier, Text]] = None,
+        description: Text = "",
     ) -> ModelTool:
         """Create a new model tool."""
         if function is not None and isinstance(function, str):
@@ -154,7 +129,7 @@ def create_model_tool(
                         break
                 if isinstance(supplier, str):
                     supplier = None
-        return ModelTool(function=function, supplier=supplier, model=model)
+        return ModelTool(function=function, supplier=supplier, model=model, description=description)
 
     @classmethod
     def create_pipeline_tool(cls, description: Text, pipeline: Union[Pipeline, Text]) -> PipelineTool:
@@ -164,37 +139,42 @@ def create_pipeline_tool(cls, description: Text, pipeline: Union[Pipeline, Text]
     @classmethod
     def list(cls) -> Dict:
         """List all agents available in the platform."""
+        from aixplain.factories.agent_factory.utils import build_agent
+
         url = urljoin(config.BACKEND_URL, "sdk/agents")
         headers = {"x-api-key": config.TEAM_API_KEY, "Content-Type": "application/json"}
 
+        resp = {}
         payload = {}
         logging.info(f"Start service for GET List Agents - {url} - {headers} - {json.dumps(payload)}")
         try:
             r = _request_with_retry("get", url, headers=headers)
             resp = r.json()
+        except Exception:
+            raise Exception("Agent Listing Error: Please contact the administrators.")
 
-            if 200 <= r.status_code < 300:
-                agents, page_total, total = [], 0, 0
-                results = resp
-                page_total = len(results)
-                total = len(results)
-                logging.info(f"Response for GET List Agents - Page Total: {page_total} / Total: {total}")
-                for agent in results:
-                    agents.append(build_agent(agent))
-                return {"results": agents, "page_total": page_total, "page_number": 0, "total": total}
-            else:
-                error_msg = "Agent Listing Error: Please contact the administrators."
-                if "message" in resp:
-                    msg = resp["message"]
-                    error_msg = f"Agent Listing Error (HTTP {r.status_code}): {msg}"
-                logging.exception(error_msg)
-                raise Exception(error_msg)
-        except Exception as e:
-            raise Exception(e)
+        if 200 <= r.status_code < 300:
+            agents, page_total, total = [], 0, 0
+            results = resp
+            page_total = len(results)
+            total = len(results)
+            logging.info(f"Response for GET List Agents - Page Total: {page_total} / Total: {total}")
+            for agent in results:
+                agents.append(build_agent(agent))
+            return {"results": agents, "page_total": page_total, "page_number": 0, "total": total}
+        else:
+            error_msg = "Agent Listing Error: Please contact the administrators."
+            if isinstance(resp, dict) and "message" in resp:
+                msg = resp["message"]
+                error_msg = f"Agent Listing Error (HTTP {r.status_code}): {msg}"
+            logging.exception(error_msg)
+            raise Exception(error_msg)
 
     @classmethod
     def get(cls, agent_id: Text, api_key: Optional[Text] = None) -> Agent:
         """Get agent by id."""
+        from aixplain.factories.agent_factory.utils import build_agent
+
         url = urljoin(config.BACKEND_URL, f"sdk/agents/{agent_id}")
         if config.AIXPLAIN_API_KEY != "":
             headers = {"x-aixplain-key": f"{config.AIXPLAIN_API_KEY}", "Content-Type": "application/json"}