Eng 812 update agents

aixplain/enums/asset_status.py

@@ -26,6 +26,7 @@
 
 
 class AssetStatus(Text, Enum):
+    DRAFT = "draft"
     HIDDEN = "hidden"
     SCHEDULED = "scheduled"
     ONBOARDING = "onboarding"

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
@@ -165,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"}

aixplain/factories/agent_factory/utils.py

@@ -12,20 +12,22 @@
 
 def build_agent(payload: Dict, api_key: Text = config.TEAM_API_KEY) -> Agent:
     """Instantiate a new agent in the platform."""
-    tools = payload["assets"]
-    for i, tool in enumerate(tools):
+    tools_dict = payload["assets"]
+    tools = []
+    for tool in tools_dict:
         if tool["type"] == "model":
-            for supplier in Supplier:
+            supplier = "aixplain"
+            for supplier_ in Supplier:
                 if tool["supplier"] is not None and tool["supplier"].lower() in [
-                    supplier.value["code"].lower(),
-                    supplier.value["name"].lower(),
+                    supplier_.value["code"].lower(),
+                    supplier_.value["name"].lower(),
                 ]:
-                    tool["supplier"] = supplier
+                    supplier = supplier_
                     break
 
             tool = ModelTool(
-                function=Function(tool["function"]) if tool["function"] is not None else None,
-                supplier=tool["supplier"],
+                function=Function(tool.get("function", None)),
+                supplier=supplier,
                 version=tool["version"],
                 model=tool["assetId"],
                 description=tool.get("description", ""),
@@ -34,37 +36,19 @@ def build_agent(payload: Dict, api_key: Text = config.TEAM_API_KEY) -> Agent:
             tool = PipelineTool(description=tool["description"], pipeline=tool["assetId"])
         else:
             raise Exception("Agent Creation Error: Tool type not supported.")
-        tools[i] = tool
+        tools.append(tool)
 
     agent = Agent(
-        id=payload["id"],
-        name=payload["name"] if "name" in payload else "",
+        id=payload["id"] if "id" in payload else "",
+        name=payload.get("name", ""),
         tools=tools,
-        description=payload["description"] if "description" in payload else "",
-        supplier=payload["teamId"] if "teamId" in payload else None,
-        version=payload["version"] if "version" in payload else None,
-        cost=payload["cost"] if "cost" in payload else None,
-        llm_id=payload["llmId"] if "llmId" in payload else GPT_4o_ID,
+        description=payload.get("description", ""),
+        supplier=payload.get("teamId", None),
+        version=payload.get("version", None),
+        cost=payload.get("cost", None),
+        llm_id=payload.get("llmId", GPT_4o_ID),
         api_key=api_key,
         status=AssetStatus(payload["status"]),
     )
     agent.url = urljoin(config.BACKEND_URL, f"sdk/agents/{agent.id}/run")
     return agent
-
-
-def validate_llm(model_id: Text) -> None:
-    from aixplain.factories.model_factory import ModelFactory
-
-    try:
-        llm = ModelFactory.get(model_id)
-        assert llm.function == Function.TEXT_GENERATION, "Large Language Model must be a text generation model."
-    except Exception:
-        raise Exception(f"Large Language Model with ID '{model_id}' not found.")
-
-
-def validate_name(name: Text) -> None:
-    import re
-
-    assert (
-        re.match("^[a-zA-Z0-9 ]*$", name) is not None
-    ), "Agent Creation Error: Agent name must not contain special characters."