merge to test

aixplain/factories/agent_factory/__init__.py

@@ -23,6 +23,7 @@
 
 import json
 import logging
+import warnings
 
 from aixplain.enums.function import Function
 from aixplain.enums.supplier import Supplier
@@ -46,6 +47,7 @@ def create(
         cls,
         name: Text,
         description: Text,
+        role: Optional[Text] = None,
         llm_id: Text = "669a63646eb56306647e1091",
         tools: List[Tool] = [],
         api_key: Text = config.TEAM_API_KEY,
@@ -54,9 +56,15 @@ def create(
     ) -> Agent:
         """Create a new agent in the platform.
 
+        Warning:
+            The 'role' parameter was recently added and serves the same purpose as 'description' did previously: set the role of the agent as a system prompt.
+            The 'description' parameter is still required and should be used to set a short summary of the agent's purpose.
+            For the next releases, the 'role' parameter will be required.
+
         Args:
             name (Text): name of the agent
             description (Text): description of the agent role.
+            role (Text): role of the agent.
             llm_id (Text, optional): aiXplain ID of the large language model to be used as agent. Defaults to "669a63646eb56306647e1091" (GPT-4o mini).
             tools (List[Tool], optional): list of tool for the agent. Defaults to [].
             api_key (Text, optional): team/user API key. Defaults to config.TEAM_API_KEY.
@@ -66,6 +74,12 @@ def create(
         Returns:
             Agent: created Agent
         """
+        warnings.warn(
+            "The 'role' parameter was recently added and serves the same purpose as 'description' did previously: set the role of the agent as a system prompt. "
+            "The 'description' parameter is still required and should be used to set a short summary of the agent's purpose. "
+            "For the next releases, the 'role' parameter will be required.",
+            UserWarning,
+        )
         from aixplain.factories.agent_factory.utils import build_agent
 
         agent = None
@@ -81,6 +95,7 @@ def create(
             "name": name,
             "assets": [tool.to_dict() for tool in tools],
             "description": description,
+            "role": role or description,
             "supplier": supplier,
             "version": version,
             "llmId": llm_id,

aixplain/factories/agent_factory/utils.py

@@ -3,7 +3,10 @@
 import aixplain.utils.config as config
 from aixplain.enums import Function, Supplier
 from aixplain.enums.asset_status import AssetStatus
-from aixplain.modules.agent import Agent, ModelTool, PipelineTool, PythonInterpreterTool
+from aixplain.modules.agent import Agent
+from aixplain.modules.agent.tool.model_tool import ModelTool
+from aixplain.modules.agent.tool.pipeline_tool import PipelineTool
+from aixplain.modules.agent.tool.python_interpreter_tool import PythonInterpreterTool
 from aixplain.modules.agent.tool.custom_python_code_tool import CustomPythonCodeTool
 from typing import Dict, Text
 from urllib.parse import urljoin
@@ -50,6 +53,7 @@ def build_agent(payload: Dict, api_key: Text = config.TEAM_API_KEY) -> Agent:
         name=payload.get("name", ""),
         tools=tools,
         description=payload.get("description", ""),
+        role=payload.get("role", ""),
         supplier=payload.get("teamId", None),
         version=payload.get("version", None),
         cost=payload.get("cost", None),

aixplain/factories/asset_factory.py

@@ -22,7 +22,7 @@
 """
 
 from abc import abstractmethod
-from typing import List, Text
+from typing import Text
 from aixplain.modules.asset import Asset
 from aixplain.utils import config
 

aixplain/factories/benchmark_factory.py

@@ -43,7 +43,6 @@ class BenchmarkFactory:
         backend_url (str): The URL for the backend.
     """
 
-
     backend_url = config.BACKEND_URL
 
     @classmethod

aixplain/factories/dataset_factory.py

@@ -57,7 +57,6 @@ class DatasetFactory(AssetFactory):
         backend_url (str): The URL for the backend.
     """
 
-
     backend_url = config.BACKEND_URL
 
     @classmethod

aixplain/factories/model_factory/__init__.py

@@ -43,8 +43,8 @@ class ModelFactory:
     @classmethod
     def create_utility_model(
         cls,
-        name: Text,
-        code: Union[Text, Callable],
+        name: Optional[Text] = None,
+        code: Union[Text, Callable] = None,
         inputs: List[UtilityModelInput] = [],
         description: Optional[Text] = None,
         output_examples: Text = "",
@@ -153,6 +153,8 @@ def list(
         sort_order: SortOrder = SortOrder.ASCENDING,
         page_number: int = 0,
         page_size: int = 20,
+        model_ids: Optional[List[Text]] = None,
+        api_key: Optional[Text] = None,
     ) -> List[Model]:
         """Gets the first k given models based on the provided task and language filters
 
@@ -165,25 +167,44 @@ def list(
             sort_by (Optional[SortBy], optional): sort the retrived models by a specific attribute,
             page_number (int, optional): page number. Defaults to 0.
             page_size (int, optional): page size. Defaults to 20.
+            model_ids (Optional[List[Text]], optional): model ids to filter. Defaults to None.
+            api_key (Optional[Text], optional): Team API key. Defaults to None.
 
         Returns:
             List[Model]: List of models based on given filters
         """
-        from aixplain.factories.model_factory.utils import get_assets_from_page
-
-        models, total = get_assets_from_page(
-            query,
-            page_number,
-            page_size,
-            function,
-            suppliers,
-            source_languages,
-            target_languages,
-            is_finetunable,
-            ownership,
-            sort_by,
-            sort_order,
-        )
+        if model_ids is not None:
+            from aixplain.factories.model_factory.utils import get_model_from_ids
+
+            assert len(model_ids) > 0, "Please provide at least one model id"
+            assert (
+                function is None
+                and suppliers is None
+                and source_languages is None
+                and target_languages is None
+                and is_finetunable is None
+                and ownership is None
+                and sort_by is None
+            ), "Cannot filter by function, suppliers, source languages, target languages, is finetunable, ownership, sort by when using model ids"
+            assert len(model_ids) <= page_size, "Page size must be greater than the number of model ids"
+            models, total = get_model_from_ids(model_ids, api_key), len(model_ids)
+        else:
+            from aixplain.factories.model_factory.utils import get_assets_from_page
+
+            models, total = get_assets_from_page(
+                query,
+                page_number,
+                page_size,
+                function,
+                suppliers,
+                source_languages,
+                target_languages,
+                is_finetunable,
+                ownership,
+                sort_by,
+                sort_order,
+                api_key,
+            )
         return {
             "results": models,
             "page_total": min(page_size, len(models)),

aixplain/factories/model_factory/utils.py

@@ -94,6 +94,7 @@ def get_assets_from_page(
     ownership: Optional[Tuple[OwnershipType, List[OwnershipType]]] = None,
     sort_by: Optional[SortBy] = None,
     sort_order: SortOrder = SortOrder.ASCENDING,
+    api_key: Optional[str] = None,
 ) -> List[Model]:
     try:
         url = urljoin(config.BACKEND_URL, "sdk/models/paginate")
@@ -146,9 +147,49 @@ def get_assets_from_page(
         all_models = resp["items"]
         from aixplain.factories.model_factory.utils import create_model_from_response
 
-        model_list = [create_model_from_response(model_info_json) for model_info_json in all_models]
+        model_list = []
+        for model_info_json in all_models:
+            model_info_json["api_key"] = config.TEAM_API_KEY
+            if api_key is not None:
+                model_info_json["api_key"] = api_key
+            model_list.append(create_model_from_response(model_info_json))
         return model_list, resp["total"]
     else:
         error_message = f"Listing Models Error: Failed to retrieve models. Status Code: {r.status_code}. Error: {resp}"
         logging.error(error_message)
         raise Exception(error_message)
+
+
+def get_model_from_ids(model_ids: List[str], api_key: Optional[str] = None) -> List[Model]:
+    from aixplain.factories.model_factory.utils import create_model_from_response
+
+    resp = None
+    try:
+        url = urljoin(config.BACKEND_URL, f"sdk/models?ids={','.join(model_ids)}")
+
+        headers = {"Authorization": f"Token {config.TEAM_API_KEY}", "Content-Type": "application/json"}
+        logging.info(f"Start service for GET Model  - {url} - {headers}")
+        r = _request_with_retry("get", url, headers=headers)
+        resp = r.json()
+
+    except Exception:
+        if resp is not None and "statusCode" in resp:
+            status_code = resp["statusCode"]
+            message = resp["message"]
+            message = f"Model Creation: Status {status_code} - {message}"
+        else:
+            message = "Model Creation: Unspecified Error"
+        logging.error(message)
+        raise Exception(f"{message}")
+    if 200 <= r.status_code < 300:
+        models = []
+        for item in resp["items"]:
+            item["api_key"] = config.TEAM_API_KEY
+            if api_key is not None:
+                item["api_key"] = api_key
+            models.append(create_model_from_response(item))
+        return models
+    else:
+        error_message = f"Model GET Error: Failed to retrieve models {model_ids}. Status Code: {r.status_code}. Error: {resp}"
+        logging.error(error_message)
+        raise Exception(error_message)

aixplain/factories/pipeline_factory/__init__.py

@@ -121,9 +121,9 @@ def get_assets_from_page(cls, page_number: int) -> List[Pipeline]:
             url = urljoin(cls.backend_url, f"sdk/pipelines/?pageNumber={page_number}")
 
             headers = {
-                    "Authorization": f"Token {config.TEAM_API_KEY}",
-                    "Content-Type": "application/json",
-                }
+                "Authorization": f"Token {config.TEAM_API_KEY}",
+                "Content-Type": "application/json",
+            }
             r = _request_with_retry("get", url, headers=headers)
             resp = r.json()
             logging.info(f"Listing Pipelines: Status of getting Pipelines on Page {page_number}: {resp}")
@@ -172,9 +172,9 @@ def list(
         url = urljoin(cls.backend_url, "sdk/pipelines/paginate")
 
         headers = {
-                "Authorization": f"Token {config.TEAM_API_KEY}",
-                "Content-Type": "application/json",
-            }
+            "Authorization": f"Token {config.TEAM_API_KEY}",
+            "Content-Type": "application/json",
+        }
 
         assert 0 < page_size <= 100, "Pipeline List Error: Page size must be greater than 0 and not exceed 100."
         payload = {

aixplain/factories/pipeline_factory/utils.py

@@ -56,13 +56,9 @@ def build_from_response(response: Dict, load_architecture: bool = False) -> Pipe
                     else:
                         node = BareAsset(asset_id=node_json["assetId"])
                 elif node_json["type"].lower() == "decision":
-                    node = Decision(
-                        routes=[Route(**route) for route in node_json["routes"]]
-                    )
+                    node = Decision(routes=[Route(**route) for route in node_json["routes"]])
                 elif node_json["type"].lower() == "router":
-                    node = Router(
-                        routes=[Route(**route) for route in node_json["routes"]]
-                    )
+                    node = Router(routes=[Route(**route) for route in node_json["routes"]])
                 elif node_json["type"].lower() == "script":
                     node = Script(
                         fileId=node_json["fileId"],
@@ -74,43 +70,21 @@ def build_from_response(response: Dict, load_architecture: bool = False) -> Pipe
                 if "inputValues" in node_json:
                     [
                         node.inputs.create_param(
-                            data_type=(
-                                DataType(input_param["dataType"])
-                                if "dataType" in input_param
-                                else None
-                            ),
+                            data_type=(DataType(input_param["dataType"]) if "dataType" in input_param else None),
                             code=input_param["code"],
-                            value=(
-                                input_param["value"] if "value" in input_param else None
-                            ),
-                            is_required=(
-                                input_param["isRequired"]
-                                if "isRequired" in input_param
-                                else False
-                            ),
+                            value=(input_param["value"] if "value" in input_param else None),
+                            is_required=(input_param["isRequired"] if "isRequired" in input_param else False),
                         )
                         for input_param in node_json["inputValues"]
                         if input_param["code"] not in node.inputs
                     ]
                 if "outputValues" in node_json:
                     [
                         node.outputs.create_param(
-                            data_type=(
-                                DataType(output_param["dataType"])
-                                if "dataType" in output_param
-                                else None
-                            ),
+                            data_type=(DataType(output_param["dataType"]) if "dataType" in output_param else None),
                             code=output_param["code"],
-                            value=(
-                                output_param["value"]
-                                if "value" in output_param
-                                else None
-                            ),
-                            is_required=(
-                                output_param["isRequired"]
-                                if "isRequired" in output_param
-                                else False
-                            ),
+                            value=(output_param["value"] if "value" in output_param else None),
+                            is_required=(output_param["isRequired"] if "isRequired" in output_param else False),
                         )
                         for output_param in node_json["outputValues"]
                         if output_param["code"] not in node.outputs