Test

.gitignore

@@ -130,3 +130,4 @@ dmypy.json
 
 # Vscode
 .vscode
+.DS_Store

aixplain/factories/__init__.py

@@ -30,3 +30,4 @@
 from .model_factory import ModelFactory
 from .pipeline_factory import PipelineFactory
 from .finetune_factory import FinetuneFactory
+from .wallet_factory import WalletFactory

aixplain/factories/agent_factory/__init__.py

@@ -34,7 +34,7 @@
 from aixplain.utils import config
 from typing import Dict, List, Optional, Text, Union
 
-from aixplain.factories.agent_factory.utils import build_agent
+from aixplain.factories.agent_factory.utils import build_agent, validate_llm
 from aixplain.utils.file_utils import _request_with_retry
 from urllib.parse import urljoin
 
@@ -50,8 +50,30 @@ def create(
         api_key: Text = config.TEAM_API_KEY,
         supplier: Union[Dict, Text, Supplier, int] = "aiXplain",
         version: Optional[Text] = None,
+        use_mentalist_and_inspector: bool = False,
     ) -> Agent:
-        """Create a new agent in the platform."""
+        """Create a new agent in the platform.
+
+        Args:
+            name (Text): name of the agent
+            llm_id (Text): aiXplain ID of the large language model to be used as agent.
+            tools (List[Tool], optional): list of tool for the agent. Defaults to [].
+            description (Text, optional): description of the agent role. Defaults to "".
+            api_key (Text, optional): team/user API key. Defaults to config.TEAM_API_KEY.
+            supplier (Union[Dict, Text, Supplier, int], optional): owner of the agent. Defaults to "aiXplain".
+            version (Optional[Text], optional): version of the agent. Defaults to None.
+            use_mentalist_and_inspector (bool, optional): flag to enable mentalist and inspector agents (which only works when a supervisor is enabled). Defaults to False.
+
+        Returns:
+            Agent: created Agent
+        """
+        # validate LLM ID
+        validate_llm(llm_id)
+
+        orchestrator_llm_id, mentalist_and_inspector_llm_id = llm_id, None
+        if use_mentalist_and_inspector is True:
+            mentalist_and_inspector_llm_id = llm_id
+
         try:
             agent = None
             url = urljoin(config.BACKEND_URL, "sdk/agents")
@@ -65,6 +87,7 @@ def create(
             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,
@@ -76,6 +99,7 @@ def create(
                         }
                     )
                 elif isinstance(tool, PipelineTool):
+                    tool.validate()
                     tool_payload.append(
                         {
                             "assetId": tool.pipeline,
@@ -92,9 +116,10 @@ def create(
                 "description": description,
                 "supplier": supplier,
                 "version": version,
+                "llmId": llm_id,
+                "supervisorId": orchestrator_llm_id,
+                "plannerId": mentalist_and_inspector_llm_id,
             }
-            if llm_id is not None:
-                payload["llmId"] = llm_id
 
             logging.info(f"Start service for POST Create Agent  - {url} - {headers} - {json.dumps(payload)}")
             r = _request_with_retry("post", url, headers=headers, json=payload)

aixplain/factories/agent_factory/utils.py

@@ -47,3 +47,13 @@ def build_agent(payload: Dict, api_key: Text = config.TEAM_API_KEY) -> Agent:
     )
     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.")

aixplain/factories/wallet_factory.py

@@ -0,0 +1,28 @@
+import aixplain.utils.config as config
+from aixplain.modules.wallet import Wallet
+from aixplain.utils.file_utils import _request_with_retry
+import logging
+from typing import Text
+
+
+class WalletFactory:
+    aixplain_key = config.AIXPLAIN_API_KEY
+    backend_url = config.BACKEND_URL
+
+    @classmethod
+    def get(cls, api_key: Text = config.TEAM_API_KEY) -> Wallet:     
+        """Get wallet information"""
+        try:
+            resp = None
+            url = f"{cls.backend_url}/sdk/billing/wallet"
+            headers = {"Authorization": f"Token {api_key}", "Content-Type": "application/json"}
+            logging.info(f"Start fetching billing information from - {url} - {headers}")
+            headers = {"Content-Type": "application/json", "x-api-key": api_key}
+            r = _request_with_retry("get", url, headers=headers)
+            resp = r.json()
+            total_balance = float(resp.get("totalBalance", 0.0))
+            reserved_balance = float(resp.get("reservedBalance", 0.0))
+
+            return Wallet(total_balance=total_balance, reserved_balance=reserved_balance)
+        except Exception as e:
+            raise Exception(f"Failed to get the wallet credit information. Error: {str(e)}")

aixplain/modules/agent/tool/model_tool.py

@@ -66,13 +66,23 @@ def __init__(
 
         if model is not None:
             if isinstance(model, Text) is True:
-                from aixplain.factories.model_factory import ModelFactory
-
-                model = ModelFactory.get(model)
+                self.model = model
+                model = self.validate()
             function = model.function
             if isinstance(model.supplier, Supplier):
                 supplier = model.supplier
             model = model.id
         self.supplier = supplier
         self.model = model
         self.function = function
+
+    def validate(self) -> Model:
+        from aixplain.factories.model_factory import ModelFactory
+
+        try:
+            model = None
+            if self.model is not None:
+                model = ModelFactory.get(self.model)
+            return model
+        except Exception:
+            raise Exception(f"Model Tool Unavailable. Make sure Model '{self.model}' exists or you have access to it.")

aixplain/modules/agent/tool/pipeline_tool.py

@@ -50,3 +50,11 @@ def __init__(
         if isinstance(pipeline, Pipeline):
             pipeline = pipeline.id
         self.pipeline = pipeline
+
+    def validate(self):
+        from aixplain.factories.pipeline_factory import PipelineFactory
+
+        try:
+            PipelineFactory.get(self.pipeline)
+        except Exception:
+            raise Exception(f"Pipeline Tool Unavailable. Make sure Pipeline '{self.pipeline}' exists or you have access to it.")

aixplain/modules/model/__init__.py

@@ -239,11 +239,27 @@ def run_async(self, data: Union[Text, Dict], name: Text = "model_process", param
 
         resp = None
         try:
-            resp = r.json()
-            logging.info(f"Result of request for {name} - {r.status_code} - {resp}")
-
-            poll_url = resp["data"]
-            response = {"status": "IN_PROGRESS", "url": poll_url}
+            if 200 <= r.status_code < 300:
+                resp = r.json()
+                logging.info(f"Result of request for {name} - {r.status_code} - {resp}")
+                poll_url = resp["data"]
+                response = {"status": "IN_PROGRESS", "url": poll_url}
+            else:
+                if r.status_code == 401:
+                    error = "Unauthorized API key: Please verify the spelling of the API key and its current validity."
+                elif 460 <= r.status_code < 470:
+                    error = "Subscription-related error: Please ensure that your subscription is active and has not expired."
+                elif 470 <= r.status_code < 480:
+                    error = "Billing-related error: Please ensure you have enough credits to run this model. "
+                elif 480 <= r.status_code < 490:
+                    error = "Supplier-related error: Please ensure that the selected supplier provides the model you are trying to access."
+                elif 490 <= r.status_code < 500:
+                    error = "Validation-related error: Please ensure all required fields are provided and correctly formatted."
+                else:
+                    status_code = str(r.status_code)
+                    error = f"Status {status_code}: Unspecified error: An unspecified error occurred while processing your request."
+                response = {"status": "FAILED", "error_message": error}
+                logging.error(f"Error in request for {name} - {r.status_code}: {error}")
         except Exception:
             response = {"status": "FAILED"}
             msg = f"Error in request for {name} - {traceback.format_exc()}"

aixplain/modules/model/llm_model.py

@@ -214,11 +214,27 @@ def run_async(
 
         resp = None
         try:
-            resp = r.json()
-            logging.info(f"Result of request for {name} - {r.status_code} - {resp}")
-
-            poll_url = resp["data"]
-            response = {"status": "IN_PROGRESS", "url": poll_url}
+            if 200 <= r.status_code < 300:
+                resp = r.json()
+                logging.info(f"Result of request for {name} - {r.status_code} - {resp}")
+                poll_url = resp["data"]
+                response = {"status": "IN_PROGRESS", "url": poll_url}
+            else:
+                if r.status_code == 401:
+                    error = "Unauthorized API key: Please verify the spelling of the API key and its current validity."
+                elif 460 <= r.status_code < 470:
+                    error = "Subscription-related error: Please ensure that your subscription is active and has not expired."
+                elif 470 <= r.status_code < 480:
+                    error = "Billing-related error: Please ensure you have enough credits to run this model. "
+                elif 480 <= r.status_code < 490:
+                    error = "Supplier-related error: Please ensure that the selected supplier provides the model you are trying to access."
+                elif 490 <= r.status_code < 500:
+                    error = "Validation-related error: Please ensure all required fields are provided and correctly formatted."
+                else:
+                    status_code = str(r.status_code)
+                    error = f"Status {status_code}: Unspecified error: An unspecified error occurred while processing your request."
+                response = {"status": "FAILED", "error_message": error}
+                logging.error(f"Error in request for {name} - {r.status_code}: {error}")               
         except Exception:
             response = {"status": "FAILED"}
             msg = f"Error in request for {name} - {traceback.format_exc()}"

aixplain/modules/pipeline/asset.py

@@ -323,11 +323,27 @@ def run_async(
 
         resp = None
         try:
-            resp = r.json()
-            logging.info(f"Result of request for {name}  - {r.status_code} - {resp}")
-
-            poll_url = resp["url"]
-            response = {"status": "IN_PROGRESS", "url": poll_url}
+            if 200 <= r.status_code < 300:
+                resp = r.json()
+                logging.info(f"Result of request for {name}  - {r.status_code} - {resp}")
+                poll_url = resp["url"]
+                response = {"status": "IN_PROGRESS", "url": poll_url}
+            else:
+                if r.status_code == 401:
+                    error = "Unauthorized API key: Please verify the spelling of the API key and its current validity."
+                elif 460 <= r.status_code < 470:
+                    error = "Subscription-related error: Please ensure that your subscription is active and has not expired."
+                elif 470 <= r.status_code < 480:
+                    error = "Billing-related error: Please ensure you have enough credits to run this pipeline. "
+                elif 480 <= r.status_code < 490:
+                    error = "Supplier-related error: Please ensure that the selected supplier provides the pipeline you are trying to access."
+                elif 490 <= r.status_code < 500:
+                    error = "Validation-related error: Please ensure all required fields are provided and correctly formatted."
+                else:
+                    status_code = str(r.status_code)
+                    error = f"Status {status_code}: Unspecified error: An unspecified error occurred while processing your request."
+                response = {"status": "FAILED", "error_message": error}
+                logging.error(f"Error in request for {name} - {r.status_code}: {error}")               
         except Exception:
             response = {"status": "FAILED"}
             if resp is not None:

aixplain/modules/wallet.py

@@ -0,0 +1,35 @@
+__author__ = "aixplain"
+
+"""
+Copyright 2024 The aiXplain SDK authors
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+
+Author: aiXplain Team
+Date: August 20th 2024
+Description:
+    Wallet Class
+"""
+
+
+class Wallet:
+    def __init__(self, total_balance: float, reserved_balance: float):
+        """
+        Args:
+            total_balance (float): total credit balance
+            reserved_balance (float): reserved credit balance
+            available_balance (float): available balance (total - credit)
+        """
+        self.total_balance = total_balance
+        self.reserved_balance = reserved_balance
+        self.available_balance =  total_balance-reserved_balance

tests/functional/agent/agent_functional_test.py

@@ -75,3 +75,9 @@ def test_list_agents():
     assert "results" in agents
     agents_result = agents["results"]
     assert type(agents_result) is list
+
+
+def test_fail_non_existent_llm():
+    with pytest.raises(Exception) as exc_info:
+        AgentFactory.create(name="Test Agent", llm_id="non_existent_llm", tools=[])
+    assert str(exc_info.value) == "Large Language Model with ID 'non_existent_llm' not found."

tests/unit/agent_test.py

@@ -2,6 +2,9 @@
 import requests_mock
 from aixplain.modules import Agent
 from aixplain.utils import config
+from aixplain.factories import AgentFactory
+from aixplain.modules.agent import PipelineTool, ModelTool
+from urllib.parse import urljoin
 
 
 def test_fail_no_data_query():
@@ -61,3 +64,71 @@ def test_sucess_query_content():
         response = agent.run_async(data={"query": "Translate the text: {{input1}}"}, content={"input1": "Hello, how are you?"})
     assert response["status"] == ref_response["status"]
     assert response["url"] == ref_response["data"]
+
+
+def test_invalid_pipelinetool():
+    with pytest.raises(Exception) as exc_info:
+        AgentFactory.create(
+            name="Test", tools=[PipelineTool(pipeline="309851793", description="Test")], llm_id="6646261c6eb563165658bbb1"
+        )
+    assert str(exc_info.value) == "Pipeline Tool Unavailable. Make sure Pipeline '309851793' exists or you have access to it."
+
+
+def test_invalid_modeltool():
+    with pytest.raises(Exception) as exc_info:
+        AgentFactory.create(name="Test", tools=[ModelTool(model="309851793")], llm_id="6646261c6eb563165658bbb1")
+    assert str(exc_info.value) == "Model Tool Unavailable. Make sure Model '309851793' exists or you have access to it."
+
+
+def test_create_agent():
+    from aixplain.enums import Supplier
+
+    with requests_mock.Mocker() as mock:
+        url = urljoin(config.BACKEND_URL, "sdk/agents")
+        headers = {"x-api-key": config.TEAM_API_KEY, "Content-Type": "application/json"}
+
+        ref_response = {
+            "id": "123",
+            "name": "Test Agent",
+            "description": "Test Agent Description",
+            "teamId": "123",
+            "version": "1.0",
+            "status": "onboarded",
+            "llmId": "6646261c6eb563165658bbb1",
+            "pricing": {"currency": "USD", "value": 0.0},
+            "assets": [
+                {
+                    "type": "model",
+                    "supplier": "openai",
+                    "version": "1.0",
+                    "assetId": "6646261c6eb563165658bbb1",
+                    "function": "text-generation",
+                }
+            ],
+        }
+        mock.post(url, headers=headers, json=ref_response)
+
+        url = urljoin(config.BACKEND_URL, "sdk/models/6646261c6eb563165658bbb1")
+        model_ref_response = {
+            "id": "6646261c6eb563165658bbb1",
+            "name": "Test LLM",
+            "description": "Test LLM Description",
+            "function": {"id": "text-generation"},
+            "supplier": "openai",
+            "version": {"id": "1.0"},
+            "status": "onboarded",
+            "pricing": {"currency": "USD", "value": 0.0},
+        }
+        mock.get(url, headers=headers, json=model_ref_response)
+
+        agent = AgentFactory.create(
+            name="Test Agent",
+            description="Test Agent Description",
+            llm_id="6646261c6eb563165658bbb1",
+            tools=[AgentFactory.create_model_tool(supplier=Supplier.OPENAI, function="text-generation")],
+        )
+
+    assert agent.name == ref_response["name"]
+    assert agent.description == ref_response["description"]
+    assert agent.llm_id == ref_response["llmId"]
+    assert agent.tools[0].function.value == ref_response["assets"][0]["function"]