Merge to test

aixplain/enums/__init__.py

@@ -20,3 +20,4 @@
 from .asset_status import AssetStatus
 from .index_stores import IndexStores
 from .function_type import FunctionType
+from .code_interpeter import CodeInterpreterModel

aixplain/enums/code_interpeter.py

@@ -0,0 +1,10 @@
+from enum import Enum
+
+
+class CodeInterpreterModel(str, Enum):
+    """Code Interpreter Model IDs"""
+
+    PYTHON_AZURE = "67476fa16eb563d00060ad62"
+
+    def __str__(self):
+        return self._value_

aixplain/enums/function_type.py

@@ -33,3 +33,4 @@ class FunctionType(Enum):
     SEARCH = "search"
     INTEGRATION = "connector"
     CONNECTION = "connection"
+    MCPSERVER = 'mcpserver'

aixplain/factories/agent_factory/__init__.py

@@ -92,13 +92,6 @@ def create(
             # Use default GPT-4o if no LLM specified
             llm = get_llm_instance("669a63646eb56306647e1091", api_key=api_key)
 
-        if instructions is None:
-            warnings.warn(
-                "Use `instructions` to define the **system prompt**. "
-                "Use `description` to provide a **short summary** of the agent for metadata and dashboard display. "
-                "Note: In upcoming releases, `instructions` will become a required parameter.",
-                UserWarning,
-            )
         warnings.warn(
             "Use `llm` to define the large language model (aixplain.modules.model.llm_model.LLM) to be used as agent. "
             "Use `llm_id` to provide the model ID of the large language model to be used as agent. "
@@ -406,4 +399,4 @@ def get(cls, agent_id: Text, api_key: Optional[Text] = None) -> Agent:
             if "message" in resp:
                 msg = resp["message"]
             error_msg = f"Agent Get Error (HTTP {r.status_code}): {msg}"
-            raise Exception(error_msg)
+            raise Exception(error_msg)

aixplain/factories/agent_factory/utils.py

@@ -177,7 +177,7 @@ def build_agent(payload: Dict, tools: List[Tool] = None, api_key: Text = config.
         name=payload.get("name", ""),
         tools=payload_tools,
         description=payload.get("description", ""),
-        instructions=payload.get("role", ""),
+        instructions=payload.get("role"),
         supplier=payload.get("teamId", None),
         version=payload.get("version", None),
         cost=payload.get("cost", None),

aixplain/factories/benchmark_factory.py

@@ -22,7 +22,7 @@
 """
 
 import logging
-from typing import Dict, List, Text
+from typing import Dict, List, Text, Any, Tuple
 import json
 from aixplain.enums.supplier import Supplier
 from aixplain.modules import Dataset, Metric, Model
@@ -150,9 +150,9 @@ def _validate_create_benchmark_payload(cls, payload):
         if len(payload["datasets"]) != 1:
             raise Exception("Please use exactly one dataset")
         if len(payload["metrics"]) == 0:
-            raise Exception("Please use exactly one metric")
-        if len(payload["model"]) == 0:
-            raise Exception("Please use exactly one model")
+            raise Exception("Please use at least one metric")
+        if len(payload["model"]) == 0 and payload.get("models", None) is None:
+            raise Exception("Please use at least one model")
         clean_metrics_info = {}
         for metric_info in payload["metrics"]:
             metric_id = metric_info["id"]
@@ -167,6 +167,31 @@ def _validate_create_benchmark_payload(cls, payload):
             {"id": metric_id, "configurations": metric_config} for metric_id, metric_config in clean_metrics_info.items()
         ]
         return payload
+
+    @classmethod
+    def _reformat_model_list(cls, model_list: List[Model]) -> Tuple[List[Any], List[Any]]:
+        """Reformat the model list to be used in the create benchmark API
+
+        Args:
+            model_list (List[Model]): List of models to be used in the benchmark
+
+        Returns:
+            Tuple[List[Any], List[Any]]: Reformatted model lists
+
+        """
+        model_list_without_parms, model_list_with_parms = [], []
+        for model in model_list:
+            if "displayName" in model.additional_info:
+                model_list_with_parms.append({"id": model.id, "displayName": model.additional_info["displayName"], "configurations": json.dumps(model.additional_info["configuration"])})
+            else:
+                model_list_without_parms.append(model.id)
+        if len(model_list_with_parms) > 0:
+            if len(model_list_without_parms) > 0:
+                raise Exception("Please provide addditional info for all models or for none of the models")
+        else:
+            model_list_with_parms = None
+        return model_list_without_parms, model_list_with_parms
+
 
     @classmethod
     def create(cls, name: str, dataset_list: List[Dataset], model_list: List[Model], metric_list: List[Metric]) -> Benchmark:
@@ -186,15 +211,18 @@ def create(cls, name: str, dataset_list: List[Dataset], model_list: List[Model],
         try:
             url = urljoin(cls.backend_url, "sdk/benchmarks")
             headers = {"Authorization": f"Token {config.TEAM_API_KEY}", "Content-Type": "application/json"}
+            model_list_without_parms, model_list_with_parms = cls._reformat_model_list(model_list)
             payload = {
                 "name": name,
                 "datasets": [dataset.id for dataset in dataset_list],
-                "model": [model.id for model in model_list],
                 "metrics": [{"id": metric.id, "configurations": metric.normalization_options} for metric in metric_list],
+                "model": model_list_without_parms,
                 "shapScores": [],
                 "humanEvaluationReport": False,
                 "automodeTraining": False,
             }
+            if model_list_with_parms is not None:
+                payload["models"] = model_list_with_parms
             clean_payload = cls._validate_create_benchmark_payload(payload)
             payload = json.dumps(clean_payload)
             r = _request_with_retry("post", url, headers=headers, data=payload)

aixplain/factories/index_factory/__init__.py

@@ -43,6 +43,11 @@ def validate_embedding_model(model_id) -> bool:
     return model.function == Function.TEXT_EMBEDDING
 
 
+def validate_embedding_model(model_id) -> bool:
+    model = ModelFactory.get(model_id)
+    return model.function == Function.TEXT_EMBEDDING
+
+
 class IndexFactory(ModelFactory, Generic[T]):
     @classmethod
     def create(

aixplain/factories/model_factory/__init__.py

@@ -30,8 +30,6 @@
 from aixplain.factories.model_factory.mixins import ModelGetterMixin, ModelListMixin
 from typing import Callable, Dict, List, Optional, Text, Union
 
-
-
 class ModelFactory(ModelGetterMixin, ModelListMixin):
     """A static class for creating and exploring Model Objects.
 

aixplain/factories/model_factory/utils.py

@@ -118,7 +118,6 @@ def create_model_from_response(response: Dict) -> Model:
         supports_streaming=response.get("supportsStreaming", False),
         status=status,
         function_type=function_type,
-
         **additional_kwargs,
     )
 

aixplain/modules/agent/__init__.py

@@ -68,7 +68,7 @@ def __init__(
         id: Text,
         name: Text,
         description: Text,
-        instructions: Text,
+        instructions: Optional[Text] = None,
         tools: List[Union[Tool, Model]] = [],
         llm_id: Text = "6646261c6eb563165658bbb1",
         llm: Optional[LLM] = None,
@@ -370,7 +370,7 @@ def to_dict(self) -> Dict:
             "name": self.name,
             "assets": [build_tool_payload(tool) for tool in self.tools],
             "description": self.description,
-            "role": self.instructions,
+            "role": self.instructions or self.description,
             "supplier": (self.supplier.value["code"] if isinstance(self.supplier, Supplier) else self.supplier),
             "version": self.version,
             "llmId": self.llm_id if self.llm is None else self.llm.id,
@@ -395,30 +395,22 @@ def delete(self) -> None:
                 "x-api-key": config.TEAM_API_KEY,
                 "Content-Type": "application/json",
             }
-            logging.debug(
-                f"Start service for DELETE Agent  - {url} - {headers}"
-            )
+            logging.debug(f"Start service for DELETE Agent  - {url} - {headers}")
             r = _request_with_retry("delete", url, headers=headers)
-            logging.debug(
-                f"Result of request for DELETE Agent - {r.status_code}"
-            )
+            logging.debug(f"Result of request for DELETE Agent - {r.status_code}")
             if r.status_code != 200:
                 raise Exception()
         except Exception:
             try:
                 response_json = r.json()
-                error_message = response_json.get('message', '').strip('{{}}')
+                error_message = response_json.get("message", "").strip("{{}}")
 
                 if r.status_code == 403 and error_message == "err.agent_is_in_use":
                     # Get team agents that use this agent
-                    from aixplain.factories.team_agent_factory import (
-                        TeamAgentFactory
-                    )
+                    from aixplain.factories.team_agent_factory import TeamAgentFactory
+
                     team_agents = TeamAgentFactory.list()["results"]
-                    using_team_agents = [
-                        ta for ta in team_agents
-                        if any(agent.id == self.id for agent in ta.agents)
-                    ]
+                    using_team_agents = [ta for ta in team_agents if any(agent.id == self.id for agent in ta.agents)]
 
                     if using_team_agents:
                         # Scenario 1: User has access to team agents
@@ -441,15 +433,9 @@ def delete(self) -> None:
                             "referencing it."
                         )
                 else:
-                    message = (
-                        f"Agent Deletion Error (HTTP {r.status_code}): "
-                        f"{error_message}."
-                    )
+                    message = f"Agent Deletion Error (HTTP {r.status_code}): " f"{error_message}."
             except ValueError:
-                message = (
-                    f"Agent Deletion Error (HTTP {r.status_code}): "
-                    "There was an error in deleting the agent."
-                )
+                message = f"Agent Deletion Error (HTTP {r.status_code}): " "There was an error in deleting the agent."
             logging.error(message)
             raise Exception(message)
 

aixplain/modules/agent/tool/custom_python_code_tool.py

@@ -25,6 +25,7 @@
 from aixplain.modules.agent.tool import Tool
 import logging
 from aixplain.enums import AssetStatus
+from aixplain.enums.code_interpeter import CodeInterpreterModel
 
 
 class CustomPythonCodeTool(Tool):
@@ -37,11 +38,13 @@ def __init__(
         super().__init__(name=name or "", description=description, **additional_info)
         self.code = code
         self.status = AssetStatus.ONBOARDED  # TODO: change to DRAFT when we have a way to onboard the tool
+        self.id = CodeInterpreterModel.PYTHON_AZURE
 
         self.validate()
 
     def to_dict(self):
         return {
+            "id": self.id,
             "name": self.name,
             "description": self.description,
             "type": "utility",

aixplain/modules/agent/utils.py

@@ -2,7 +2,9 @@
 import re
 
 
-def process_variables(query: Union[Text, Dict], data: Union[Dict, Text], parameters: Dict, agent_description: Text) -> Text:
+def process_variables(
+    query: Union[Text, Dict], data: Union[Dict, Text], parameters: Dict, agent_description: Union[Text, None]
+) -> Text:
     from aixplain.factories.file_factory import FileFactory
 
     if isinstance(query, dict):
@@ -13,7 +15,7 @@ def process_variables(query: Union[Text, Dict], data: Union[Dict, Text], paramet
     else:
         input_data = {"input": FileFactory.to_link(query)}
 
-    variables = re.findall(r"(?<!{){([^}]+)}(?!})", agent_description)
+    variables = re.findall(r"(?<!{){([^}]+)}(?!})", agent_description or "")
     for variable in variables:
         if isinstance(data, dict):
             assert (

aixplain/modules/benchmark_job.py

@@ -3,6 +3,7 @@
 from aixplain.utils import config
 from urllib.parse import urljoin
 import pandas as pd
+import json
 from pathlib import Path
 from aixplain.utils.request_utils import _request_with_retry
 from aixplain.utils.file_utils import save_file
@@ -109,6 +110,10 @@ def get_scores(self, return_simplified=True, return_as_dataframe=True):
             scores = {}
             for iteration_info in iterations:
                 model_id = iteration_info["pipeline"]
+                pipeline_json = json.loads(iteration_info["pipelineJson"])
+                if "benchmark" in pipeline_json:
+                    model_id = pipeline_json["benchmark"]["displayName"]
+
                 model_info = {
                     "creditsUsed": round(iteration_info.get("credits", 0), 5),
                     "timeSpent": round(iteration_info.get("runtime", 0), 2),

aixplain/modules/model/__init__.py

@@ -432,6 +432,16 @@ def delete(self) -> None:
             message = "Model Deletion Error: Make sure the model exists and you are the owner."
             logging.error(message)
             raise Exception(f"{message}")
+
+    def add_additional_info_for_benchmark(self, display_name: str, configuration: Dict) -> None:
+        """Add additional info for benchmark
+
+        Args:
+            display_name (str): display name of the model
+            configuration (Dict): configuration of the model
+        """
+        self.additional_info["displayName"] = display_name
+        self.additional_info["configuration"] = configuration
 
     @classmethod
     def from_dict(cls, data: Dict) -> "Model":
@@ -451,3 +461,4 @@ def from_dict(cls, data: Dict) -> "Model":
             model_params=data.get("model_params"),
             **data.get("additional_info", {}),
         )
+

aixplain/modules/model/index_model.py

@@ -7,7 +7,6 @@
 from enum import Enum
 from typing import List
 from aixplain.enums.splitting_options import SplittingOptions
-
 import os
 
 from urllib.parse import urljoin
@@ -56,8 +55,6 @@ def __init__(
         self.split_length = split_length
         self.split_overlap = split_overlap
 
-
-
 class IndexModel(Model):
     def __init__(
         self,
@@ -125,7 +122,6 @@ def to_dict(self) -> Dict:
         data["collection_type"] = self.version.split("-", 1)[0]
         return data
 
-
     def search(self, query: str, top_k: int = 10, filters: List[IndexFilter] = []) -> ModelResponse:
         """Search for documents in the index
 

aixplain/modules/pipeline/default.py

@@ -15,5 +15,4 @@ def save(self, *args, **kwargs):
 
     def to_dict(self) -> dict:
         return self.serialize()
-
-
+

aixplain/modules/pipeline/designer/enums.py

@@ -1,5 +1,5 @@
 from enum import Enum
-
+from aixplain.enums import FunctionType
 
 class RouteType(str, Enum):
     CHECK_TYPE = "checkType"
@@ -29,15 +29,6 @@ class NodeType(str, Enum):
 class AssetType(str, Enum):
     MODEL = "MODEL"
 
-
-class FunctionType(str, Enum):
-    AI = "ai"
-    SEGMENTOR = "segmentor"
-    RECONSTRUCTOR = "reconstructor"
-    UTILITY = "utility"
-    METRIC = "metric"
-
-
 class ParamType:
     INPUT = "INPUT"
     OUTPUT = "OUTPUT"

aixplain/modules/pipeline/designer/nodes.py

@@ -1,4 +1,5 @@
 from typing import List, Union, Type, TYPE_CHECKING, Optional
+from enum import Enum
 
 from aixplain.modules import Model
 from aixplain.enums import DataType, Function
@@ -142,7 +143,11 @@ def serialize(self) -> dict:
         obj["supplier"] = self.supplier
         obj["version"] = self.version
         obj["assetType"] = self.assetType
-        obj["functionType"] = self.functionType
+        # Handle functionType as enum or string
+        if isinstance(self.functionType, Enum):
+            obj["functionType"] = self.functionType.value
+        else:
+            obj["functionType"] = self.functionType
         obj["type"] = self.type
         return obj