Merge to test

.pre-commit-config.yaml

@@ -21,5 +21,5 @@ repos:
     rev: v2.0.0 # Use the latest version
     hooks:
       - id: flake8
-        args: # arguments to configure black
-          - --ignore=E402,E501
+        args: # arguments to configure flake8
+          - --ignore=E402,E501,E203

aixplain/enums/__init__.py

@@ -1,3 +1,4 @@
+# flake8: noqa: F401 // to ignore the F401 (unused import)
 from .data_split import DataSplit
 from .data_subtype import DataSubtype
 from .data_type import DataType
@@ -14,3 +15,4 @@
 from .sort_by import SortBy
 from .sort_order import SortOrder
 from .response_status import ResponseStatus
+from .database_source import DatabaseSourceType

aixplain/enums/database_source.py

@@ -0,0 +1,47 @@
+__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: Lucas Pavanelli and Thiago Castro Ferreira and Ahmet Gunduz
+Date: March 7th 2025
+Description:
+    Database Source Type Enum
+"""
+
+from enum import Enum
+
+
+class DatabaseSourceType(Enum):
+    """Enum for database source types"""
+
+    POSTGRESQL = "postgresql"
+    SQLITE = "sqlite"
+    CSV = "csv"
+
+    @classmethod
+    def from_string(cls, source_type: str) -> "DatabaseSourceType":
+        """Convert string to DatabaseSourceType enum
+
+        Args:
+            source_type (str): Source type string
+
+        Returns:
+            DatabaseSourceType: Corresponding enum value
+        """
+        try:
+            return cls[source_type.upper()]
+        except KeyError:
+            raise ValueError(f"Invalid source type: {source_type}")

aixplain/enums/status.py

@@ -0,0 +1,8 @@
+from enum import Enum
+from typing import Text
+
+
+class Status(Text, Enum):
+    FAILED = "failed"
+    IN_PROGRESS = "in_progress"
+    SUCCESS = "success"

aixplain/factories/agent_factory/__init__.py

@@ -24,6 +24,7 @@
 import json
 import logging
 import warnings
+import os
 
 from aixplain.enums.function import Function
 from aixplain.enums.supplier import Supplier
@@ -32,7 +33,9 @@
 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 aixplain.modules.agent.tool.sql_tool import SQLTool
+from aixplain.modules.agent.tool.sql_tool import (
+    SQLTool,
+)
 from aixplain.modules.model import Model
 from aixplain.modules.pipeline import Pipeline
 from aixplain.utils import config
@@ -121,8 +124,8 @@ def create(
             "status": "draft",
             "tasks": [task.to_dict() for task in tasks],
         }
-        agent = build_agent(payload=payload, api_key=api_key)
-        agent.validate()
+        agent = build_agent(payload=payload, tools=tools, api_key=api_key)
+        agent.validate(raise_exception=True)
         response = "Unspecified error"
         try:
             logging.debug(f"Start service for POST Create Agent  - {url} - {headers} - {json.dumps(agent.to_dict())}")
@@ -132,7 +135,7 @@ def create(
             raise Exception("Agent Onboarding Error: Please contact the administrators.")
 
         if 200 <= r.status_code < 300:
-            agent = build_agent(payload=response, api_key=api_key)
+            agent = build_agent(payload=response, tools=tools, api_key=api_key)
         else:
             error_msg = f"Agent Onboarding Error: {response}"
             if "message" in response:
@@ -193,7 +196,8 @@ def create_custom_python_code_tool(cls, code: Union[Text, Callable], description
     def create_sql_tool(
         cls,
         description: Text,
-        database: Text,
+        source: str,
+        source_type: str,
         schema: Optional[Text] = None,
         tables: Optional[List[Text]] = None,
         enable_commit: bool = False,
@@ -202,15 +206,114 @@ def create_sql_tool(
 
         Args:
             description (Text): description of the database tool
-            database (Text): URL/local path of the SQLite database file
-            schema (Optional[Text], optional): database schema description (optional)
+            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)
+            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)
-
         Returns:
             SQLTool: created SQLTool
+
+        Examples:
+            # CSV - Simple
+            sql_tool = AgentFactory.create_sql_tool(
+                description="My CSV Tool",
+                source="/path/to/data.csv",
+                source_type="csv",
+                tables=["data"]
+            )
+
+            # SQLite - Simple
+            sql_tool = AgentFactory.create_sql_tool(
+                description="My SQLite Tool",
+                source="/path/to/database.sqlite",
+                source_type="sqlite",
+                tables=["users", "products"]
+            )
         """
-        return SQLTool(description=description, database=database, schema=schema, tables=tables, enable_commit=enable_commit)
+        from aixplain.modules.agent.tool.sql_tool import (
+            SQLToolError,
+            create_database_from_csv,
+            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))
+
+        database_path = None  # Final database path to pass to SQLTool
+
+        # Handle CSV source type
+        if source_type == DatabaseSourceType.CSV:
+            if not os.path.exists(source):
+                raise SQLToolError(f"CSV file '{source}' does not exist")
+            if not source.endswith(".csv"):
+                raise SQLToolError(f"File '{source}' is not a CSV file")
+
+            # 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")
+
+            try:
+                # Create database from CSV
+                schema = create_database_from_csv(source, db_path)
+                database_path = db_path
+
+                # Get table names if not provided
+                if not tables:
+                    tables = get_table_names_from_schema(schema)
+
+            except Exception as e:
+                if os.path.exists(db_path):
+                    try:
+                        os.remove(db_path)
+                    except Exception as cleanup_error:
+                        warnings.warn(f"Failed to remove temporary database file '{db_path}': {str(cleanup_error)}")
+                raise SQLToolError(f"Failed to create database from CSV: {str(e)}")
+
+        # Handle SQLite source type
+        elif source_type == DatabaseSourceType.SQLITE:
+            if not os.path.exists(source):
+                raise SQLToolError(f"Database '{source}' does not exist")
+            if not source.endswith(".db") and not source.endswith(".sqlite"):
+                raise SQLToolError(f"Database '{source}' must have .db or .sqlite extension")
+
+            database_path = source
+
+            # Infer schema from database if not provided
+            if not schema:
+                try:
+                    schema = get_table_schema(database_path)
+                except Exception as e:
+                    raise SQLToolError(f"Failed to get database schema: {str(e)}")
+
+            # Get table names if not provided
+            if not tables:
+                try:
+                    tables = get_table_names_from_schema(schema)
+                except Exception as e:
+                    raise SQLToolError(f"Failed to get table names: {str(e)}")
+
+        elif source_type == DatabaseSourceType.POSTGRESQL:
+            raise SQLToolError("PostgreSQL is not supported yet")
+
+        # Create and return SQLTool
+        return SQLTool(
+            description=description,
+            database=database_path,
+            schema=schema,
+            tables=tables,
+            enable_commit=enable_commit,
+        )
 
     @classmethod
     def list(cls) -> Dict:

aixplain/factories/agent_factory/utils.py

@@ -1,69 +1,93 @@
 __author__ = "thiagocastroferreira"
 
+import logging
 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
+from aixplain.modules.agent.tool import Tool
 from aixplain.modules.agent.agent_task import AgentTask
 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 aixplain.modules.agent.tool.sql_tool import SQLTool
-from typing import Dict, Text
+from typing import Dict, Text, List
 from urllib.parse import urljoin
 
 GPT_4o_ID = "6646261c6eb563165658bbb1"
 
 
-def build_agent(payload: Dict, api_key: Text = config.TEAM_API_KEY) -> Agent:
+def build_tool(tool: Dict):
+    """Build a tool from a dictionary.
+
+    Args:
+        tool (Dict): Tool dictionary.
+
+    Returns:
+        Tool: Tool object.
+    """
+    if tool["type"] == "model":
+        supplier = "aixplain"
+        for supplier_ in Supplier:
+            if isinstance(tool["supplier"], str):
+                if tool["supplier"] is not None and tool["supplier"].lower() in [
+                    supplier_.value["code"].lower(),
+                    supplier_.value["name"].lower(),
+                ]:
+                    supplier = supplier_
+                    break
+        tool = ModelTool(
+            function=Function(tool.get("function", None)),
+            supplier=supplier,
+            version=tool["version"],
+            model=tool["assetId"],
+            description=tool.get("description", ""),
+            parameters=tool.get("parameters", None),
+        )
+    elif tool["type"] == "pipeline":
+        tool = PipelineTool(description=tool["description"], pipeline=tool["assetId"])
+    elif tool["type"] == "utility":
+        if tool.get("utilityCode", None) is not None:
+            tool = CustomPythonCodeTool(description=tool["description"], code=tool["utilityCode"])
+        else:
+            tool = PythonInterpreterTool()
+    elif tool["type"] == "sql":
+        parameters = {parameter["name"]: parameter["value"] for parameter in tool.get("parameters", [])}
+        database = parameters.get("database")
+        schema = parameters.get("schema")
+        tables = parameters.get("tables", None)
+        tables = tables.split(",") if tables is not None else None
+        enable_commit = parameters.get("enable_commit", False)
+        tool = SQLTool(
+            description=tool["description"], database=database, schema=schema, tables=tables, enable_commit=enable_commit
+        )
+    else:
+        raise Exception("Agent Creation Error: Tool type not supported.")
+
+    return tool
+
+
+def build_agent(payload: Dict, tools: List[Tool] = None, api_key: Text = config.TEAM_API_KEY) -> Agent:
     """Instantiate a new agent in the platform."""
     tools_dict = payload["assets"]
-    tools = []
-    for tool in tools_dict:
-        if tool["type"] == "model":
-            supplier = "aixplain"
-            for supplier_ in Supplier:
-                if isinstance(tool["supplier"], str):
-                    if tool["supplier"] is not None and tool["supplier"].lower() in [
-                        supplier_.value["code"].lower(),
-                        supplier_.value["name"].lower(),
-                    ]:
-                        supplier = supplier_
-                        break
-            tool = ModelTool(
-                function=Function(tool.get("function", None)),
-                supplier=supplier,
-                version=tool["version"],
-                model=tool["assetId"],
-                description=tool.get("description", ""),
-                parameters=tool.get("parameters", None),
-            )
-        elif tool["type"] == "pipeline":
-            tool = PipelineTool(description=tool["description"], pipeline=tool["assetId"])
-        elif tool["type"] == "utility":
-            if tool.get("utilityCode", None) is not None:
-                tool = CustomPythonCodeTool(description=tool["description"], code=tool["utilityCode"])
-            else:
-                tool = PythonInterpreterTool()
-        elif tool["type"] == "sql":
-            parameters = {parameter["name"]: parameter["value"] for parameter in tool.get("parameters", [])}
-            database = parameters.get("database")
-            schema = parameters.get("schema")
-            tables = parameters.get("tables", None)
-            tables = tables.split(",") if tables is not None else None
-            enable_commit = parameters.get("enable_commit", False)
-            tool = SQLTool(
-                description=tool["description"], database=database, schema=schema, tables=tables, enable_commit=enable_commit
-            )
-        else:
-            raise Exception("Agent Creation Error: Tool type not supported.")
-        tools.append(tool)
+    payload_tools = tools
+    if payload_tools is None:
+        payload_tools = []
+        for tool in tools_dict:
+            try:
+                payload_tools.append(build_tool(tool))
+            except Exception:
+                logging.warning(
+                    f"Tool {tool['assetId']} is not available. Make sure it exists or you have access to it. "
+                    "If you think this is an error, please contact the administrators."
+                )
+                continue
 
     agent = Agent(
         id=payload["id"] if "id" in payload else "",
         name=payload.get("name", ""),
-        tools=tools,
+        tools=payload_tools,
         description=payload.get("description", ""),
         instructions=payload.get("role", ""),
         supplier=payload.get("teamId", None),