Merge to test

aixplain/enums/__init__.py

@@ -16,3 +16,5 @@
 from .sort_order import SortOrder
 from .response_status import ResponseStatus
 from .database_source import DatabaseSourceType
+from .embedding_model import EmbeddingModel
+from .asset_status import AssetStatus

aixplain/enums/embedding_model.py

@@ -0,0 +1,30 @@
+__author__ = "aiXplain"
+
+"""
+Copyright 2023 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: February 17th 2025
+Description:
+    Embedding Model Enum
+"""
+
+from enum import Enum
+
+
+class EmbeddingModel(Enum):
+    SNOWFLAKE_ARCTIC_EMBED_M_LONG = "6658d40729985c2cf72f42ec"
+    OPENAI_ADA002 = "6734c55df127847059324d9e"
+    SNOWFLAKE_ARCTIC_EMBED_L_V2_0 = "678a4f8547f687504744960a"
+    JINA_CLIP_V2_MULTIMODAL = "67c5f705d8f6a65d6f74d732"
+
+    def __str__(self):
+        return self._value_

aixplain/factories/agent_factory/__init__.py

@@ -43,6 +43,7 @@
 
 from aixplain.utils.file_utils import _request_with_retry
 from urllib.parse import urljoin
+from aixplain.enums import DatabaseSourceType
 
 
 class AgentFactory:
@@ -197,7 +198,7 @@ def create_sql_tool(
         cls,
         description: Text,
         source: str,
-        source_type: str,
+        source_type: Union[str, DatabaseSourceType],
         schema: Optional[Text] = None,
         tables: Optional[List[Text]] = None,
         enable_commit: bool = False,
@@ -207,7 +208,7 @@ def create_sql_tool(
         Args:
             description (Text): description of the database tool
             source (Union[Text, Dict]): database source - can be a connection string or dictionary with connection details
-            source_type (Text): type of source (postgresql, sqlite, csv)
+            source_type (Union[str, DatabaseSourceType]): type of source (postgresql, sqlite, csv) or DatabaseSourceType enum
             schema (Optional[Text], optional): database schema description
             tables (Optional[List[Text]], optional): table names to work with (optional)
             enable_commit (bool, optional): enable to modify the database (optional)
@@ -237,18 +238,23 @@ def create_sql_tool(
             get_table_schema,
             get_table_names_from_schema,
         )
-        from aixplain.enums import DatabaseSourceType
 
         if not source:
             raise SQLToolError("Source must be provided")
         if not source_type:
             raise SQLToolError("Source type must be provided")
 
         # Validate source type
-        try:
-            source_type = DatabaseSourceType.from_string(source_type)
-        except ValueError as e:
-            raise SQLToolError(str(e))
+        if isinstance(source_type, str):
+            try:
+                source_type = DatabaseSourceType.from_string(source_type)
+            except ValueError as e:
+                raise SQLToolError(str(e))
+        elif isinstance(source_type, DatabaseSourceType):
+            # Already the correct type, no conversion needed
+            pass
+        else:
+            raise SQLToolError(f"Source type must be either a string or DatabaseSourceType enum, got {type(source_type)}")
 
         database_path = None  # Final database path to pass to SQLTool
 
@@ -258,14 +264,17 @@ def create_sql_tool(
                 raise SQLToolError(f"CSV file '{source}' does not exist")
             if not source.endswith(".csv"):
                 raise SQLToolError(f"File '{source}' is not a CSV file")
+            if tables and len(tables) > 1:
+                raise SQLToolError("CSV source type only supports one table")
 
             # Create database name from CSV filename or use custom table name
             base_name = os.path.splitext(os.path.basename(source))[0]
             db_path = os.path.join(os.path.dirname(source), f"{base_name}.db")
+            table_name = tables[0] if tables else None
 
             try:
                 # Create database from CSV
-                schema = create_database_from_csv(source, db_path)
+                schema = create_database_from_csv(source, db_path, table_name)
                 database_path = db_path
 
                 # Get table names if not provided

aixplain/factories/index_factory.py

@@ -1,16 +1,20 @@
 from aixplain.modules.model.index_model import IndexModel
 from aixplain.factories import ModelFactory
-from aixplain.enums import Function, ResponseStatus, SortBy, SortOrder, OwnershipType, Supplier
+from aixplain.enums import EmbeddingModel, Function, ResponseStatus, SortBy, SortOrder, OwnershipType, Supplier
 from typing import Optional, Text, Union, List, Tuple
 
+AIR_MODEL_ID = "66eae6656eb56311f2595011"
+
 
 class IndexFactory(ModelFactory):
     @classmethod
-    def create(cls, name: Text, description: Text) -> IndexModel:
+    def create(
+        cls, name: Text, description: Text, embedding_model: EmbeddingModel = EmbeddingModel.OPENAI_ADA002
+    ) -> IndexModel:
         """Create a new index collection"""
-        model = cls.get("66eae6656eb56311f2595011")
+        model = cls.get(AIR_MODEL_ID)
 
-        data = {"data": name, "description": description}
+        data = {"data": name, "description": description, "model": embedding_model.value}
         response = model.run(data=data)
         if response.status == ResponseStatus.SUCCESS:
             model_id = response.data
@@ -19,7 +23,7 @@ def create(cls, name: Text, description: Text) -> IndexModel:
 
         error_message = f"Index Factory Exception: {response.error_message}"
         if error_message == "":
-            error_message = "Index Factory Exception:An error occurred while creating the index collection."
+            error_message = "Index Factory Exception: An error occurred while creating the index collection."
         raise Exception(error_message)
 
     @classmethod

aixplain/factories/team_agent_factory/__init__.py

@@ -23,15 +23,15 @@
 
 import json
 import logging
+from typing import Dict, List, Optional, Text, Union
+from urllib.parse import urljoin
 
 from aixplain.enums.supplier import Supplier
 from aixplain.modules.agent import Agent
-from aixplain.modules.team_agent import TeamAgent
+from aixplain.modules.team_agent import TeamAgent, InspectorTarget
 from aixplain.utils import config
 from aixplain.factories.team_agent_factory.utils import build_team_agent
 from aixplain.utils.file_utils import _request_with_retry
-from typing import Dict, List, Optional, Text, Union
-from urllib.parse import urljoin
 
 
 class TeamAgentFactory:
@@ -47,9 +47,29 @@ def create(
         version: Optional[Text] = None,
         use_mentalist: bool = True,
         use_inspector: bool = True,
+        num_inspectors: int = 1,
+        inspector_targets: List[Union[InspectorTarget, Text]] = [InspectorTarget.STEPS],
         use_mentalist_and_inspector: bool = False,  # TODO: remove this
     ) -> TeamAgent:
-        """Create a new team agent in the platform."""
+        """Create a new team agent in the platform.
+
+        Args:
+            name: The name of the team agent.
+            agents: A list of agents to be added to the team.
+            llm_id: The ID of the LLM to be used for the team agent.
+            description: The description of the team agent.
+            api_key: The API key to be used for the team agent.
+            supplier: The supplier of the team agent.
+            version: The version of the team agent.
+            use_mentalist: Whether to use the mentalist agent.
+            use_inspector: Whether to use the inspector agent.
+            num_inspectors: The number of inspectors to be used for each inspection.
+            inspector_targets: Which stages to be inspected during an execution of the team agent. (steps, output)
+            use_mentalist_and_inspector: Whether to use the mentalist and inspector agents. (legacy)
+
+        Returns:
+            A new team agent instance.
+        """
         assert len(agents) > 0, "TeamAgent Onboarding Error: At least one agent must be provided."
         agent_list = []
         for agent in agents:
@@ -68,8 +88,22 @@ def create(
                 assert isinstance(agent, Agent), "TeamAgent Onboarding Error: Agents must be instances of Agent class"
             agent_list.append(agent_obj)
 
-        if use_inspector and not use_mentalist:
-            raise Exception("TeamAgent Onboarding Error: To use the Inspector agent, you must enable Mentalist.")
+        # NOTE: backend expects max_inspectors (for "generated" inspectors)
+        max_inspectors = num_inspectors
+
+        if use_inspector:
+            try:
+                # convert to enum if string and check its validity
+                inspector_targets = [InspectorTarget(target) for target in inspector_targets]
+            except ValueError:
+                raise ValueError("TeamAgent Onboarding Error: Invalid inspector target. Valid targets are: steps, output")
+
+            if not use_mentalist:
+                raise Exception("TeamAgent Onboarding Error: To use the Inspector agent, you must enable Mentalist.")
+            if max_inspectors < 1:
+                raise Exception(
+                    "TeamAgent Onboarding Error: The number of inspectors must be greater than 0 when using the Inspector agent."
+                )
 
         if use_mentalist_and_inspector:
             mentalist_llm_id = llm_id
@@ -100,6 +134,8 @@ def create(
             "supervisorId": llm_id,
             "plannerId": mentalist_llm_id,
             "inspectorId": inspector_llm_id,
+            "maxInspectors": max_inspectors,
+            "inspectorTargets": inspector_targets if use_inspector else [],
             "supplier": supplier,
             "version": version,
             "status": "draft",

aixplain/factories/team_agent_factory/utils.py

@@ -1,12 +1,14 @@
 __author__ = "lucaspavanelli"
 
 import logging
+from typing import Dict, Text, List
+from urllib.parse import urljoin
+
 import aixplain.utils.config as config
 from aixplain.enums.asset_status import AssetStatus
 from aixplain.modules.agent import Agent
-from aixplain.modules.team_agent import TeamAgent
-from typing import Dict, Text, List
-from urllib.parse import urljoin
+from aixplain.modules.team_agent import TeamAgent, InspectorTarget
+
 
 GPT_4o_ID = "6646261c6eb563165658bbb1"
 
@@ -29,6 +31,8 @@ def build_team_agent(payload: Dict, agents: List[Agent] = None, api_key: Text =
                 )
                 continue
 
+    inspector_targets = [InspectorTarget(target.lower()) for target in payload.get("inspectorTargets", [])]
+
     team_agent = TeamAgent(
         id=payload.get("id", ""),
         name=payload.get("name", ""),
@@ -40,6 +44,8 @@ def build_team_agent(payload: Dict, agents: List[Agent] = None, api_key: Text =
         llm_id=payload.get("llmId", GPT_4o_ID),
         use_mentalist=True if payload.get("plannerId", None) is not None else False,
         use_inspector=True if payload.get("inspectorId", None) is not None else False,
+        max_inspectors=payload.get("maxInspectors", 1),
+        inspector_targets=inspector_targets,
         api_key=api_key,
         status=AssetStatus(payload["status"]),
     )

aixplain/modules/agent/tool/custom_python_code_tool.py

@@ -36,6 +36,7 @@ def __init__(self, code: Union[Text, Callable], description: Text = "", **additi
 
     def to_dict(self):
         return {
+            "name": self.name,
             "description": self.description,
             "type": "utility",
             "utility": "custom_python_code",

aixplain/modules/agent/tool/model_tool.py

@@ -42,6 +42,7 @@ def __init__(
         function: Optional[Union[Function, Text]] = None,
         supplier: Optional[Union[Dict, Supplier]] = None,
         model: Optional[Union[Text, Model]] = None,
+        name: Optional[Text] = None,
         description: Text = "",
         parameters: Optional[Dict] = None,
         **additional_info,
@@ -59,7 +60,8 @@ def __init__(
             function is not None or model is not None
         ), "Agent Creation Error: Either function or model must be provided when instantiating a tool."
 
-        super().__init__(name="", description=description, **additional_info)
+        name = name or ""
+        super().__init__(name=name, description=description, **additional_info)
         if function is not None:
             if isinstance(function, str):
                 function = Function(function)
@@ -104,6 +106,7 @@ def to_dict(self) -> Dict:
         return {
             "function": self.function.value if self.function is not None else None,
             "type": "model",
+            "name": self.name,
             "description": self.description,
             "supplier": supplier,
             "version": self.version if self.version else None,

aixplain/modules/agent/tool/pipeline_tool.py

@@ -20,7 +20,7 @@
 Description:
     Agentification Class
 """
-from typing import Text, Union
+from typing import Text, Union, Optional
 
 from aixplain.modules.agent.tool import Tool
 from aixplain.modules.pipeline import Pipeline
@@ -38,6 +38,7 @@ def __init__(
         self,
         description: Text,
         pipeline: Union[Text, Pipeline],
+        name: Optional[Text] = None,
         **additional_info,
     ) -> None:
         """Specialized software or resource designed to assist the AI in executing specific tasks or functions based on user commands.
@@ -46,14 +47,17 @@ def __init__(
             description (Text): description of the tool
             pipeline (Union[Text, Pipeline]): pipeline
         """
-        super().__init__("", description, **additional_info)
+        name = name or ""
+        super().__init__(name=name, description=description, **additional_info)
+
         if isinstance(pipeline, Pipeline):
             pipeline = pipeline.id
         self.pipeline = pipeline
 
     def to_dict(self):
         return {
             "assetId": self.pipeline,
+            "name": self.name,
             "description": self.description,
             "type": "pipeline",
         }

aixplain/modules/agent/tool/sql_tool.py

@@ -127,7 +127,7 @@ def get_table_schema(database_path: str) -> str:
         raise DatabaseError(f"Unexpected error while getting table schema: {str(e)}")
 
 
-def create_database_from_csv(csv_path: str, database_path: str) -> str:
+def create_database_from_csv(csv_path: str, database_path: str, table_name: str = None) -> str:
     """Create SQLite database from CSV file and return the schema"""
     if not os.path.exists(csv_path):
         raise CSVError(f"CSV file '{csv_path}' does not exist")
@@ -161,7 +161,9 @@ def create_database_from_csv(csv_path: str, database_path: str) -> str:
                 cursor = conn.cursor()
 
                 # Check if table already exists
-                table_name = clean_column_name(os.path.splitext(os.path.basename(csv_path))[0])
+
+                if table_name is None:
+                    table_name = clean_column_name(os.path.splitext(os.path.basename(csv_path))[0])
                 cursor.execute(f"SELECT name FROM sqlite_master WHERE type='table' AND name='{table_name}'")
                 if cursor.fetchone():
                     warnings.warn(f"Table '{table_name}' already exists in the database and will be replaced")
@@ -262,6 +264,7 @@ def __init__(
         schema: Optional[Text] = None,
         tables: Optional[Union[List[Text], Text]] = None,
         enable_commit: bool = False,
+        name: Optional[Text] = None,
         **additional_info,
     ) -> None:
         """Tool to execute SQL query commands in an SQLite database.
@@ -273,14 +276,18 @@ def __init__(
             tables (Optional[Union[List[Text], Text]]): table names to work with (optional)
             enable_commit (bool): enable to modify the database (optional)
         """
-        super().__init__("", description, **additional_info)
+
+        name = name or ""
+        super().__init__(name=name, description=description, **additional_info)
+
         self.database = database
         self.schema = schema
         self.tables = tables if isinstance(tables, list) else [tables] if tables else None
         self.enable_commit = enable_commit
 
     def to_dict(self) -> Dict[str, Text]:
         return {
+            "name": self.name,
             "description": self.description,
             "parameters": [
                 {"name": "database", "value": self.database},