alita-sdk 0.3.257__py3-none-any.whl → 0.3.584__py3-none-any.whl

alita_sdk/tools/sql/__init__.py

@@ -1,12 +1,15 @@
 from typing import List, Literal, Optional
 
 from langchain_core.tools import BaseToolkit, BaseTool
-from pydantic import create_model, BaseModel, ConfigDict, Field, SecretStr
+from pydantic import create_model, BaseModel, ConfigDict, Field
 
 from .api_wrapper import SQLApiWrapper
 from ..base.tool import BaseAction
 from .models import SQLDialect
-from ..utils import TOOLKIT_SPLITTER, clean_string, get_max_toolkit_length
+from ..elitea_base import filter_missconfigured_index_tools
+from ..utils import clean_string, get_max_toolkit_length
+from ...configurations.sql import SqlConfiguration
+from ...runtime.utils.constants import TOOLKIT_NAME_META, TOOL_NAME_META, TOOLKIT_TYPE_META
 
 name = "sql"
 
@@ -14,32 +17,24 @@ def get_tools(tool):
     return SQLToolkit().get_toolkit(
         selected_tools=tool['settings'].get('selected_tools', []),
         dialect=tool['settings']['dialect'],
-        host=tool['settings']['host'],
-        port=tool['settings']['port'],
-        username=tool['settings']['username'],
-        password=tool['settings']['password'],
         database_name=tool['settings']['database_name'],
+        sql_configuration=tool['settings']['sql_configuration'],
         toolkit_name=tool.get('toolkit_name')
     ).get_tools()
 
 
 class SQLToolkit(BaseToolkit):
     tools: list[BaseTool] = []
-    toolkit_max_length: int = 0
 
     @staticmethod
     def toolkit_config_schema() -> BaseModel:
         selected_tools = {x['name']: x['args_schema'].schema() for x in SQLApiWrapper.model_construct().get_available_tools()}
-        SQLToolkit.toolkit_max_length = get_max_toolkit_length(selected_tools)
         supported_dialects = (d.value for d in SQLDialect)
         return create_model(
             name,
             dialect=(Literal[tuple(supported_dialects)], Field(description="Database dialect (mysql or postgres)")),
-            host=(str, Field(description="Database server address")),
-            port=(str, Field(description="Database server port")),
-            username=(str, Field(description="Database username")),
-            password=(SecretStr, Field(description="Database password", json_schema_extra={'secret': True})),
-            database_name=(str, Field(description="Database name", json_schema_extra={'toolkit_name': True, 'max_toolkit_length': SQLToolkit.toolkit_max_length})),
+            database_name=(str, Field(description="Database name")),
+            sql_configuration=(SqlConfiguration, Field(description="SQL Configuration", json_schema_extra={'configuration_types': ['sql']})),
             selected_tools=(List[Literal[tuple(selected_tools)]], Field(default=[], json_schema_extra={'args_schemas': selected_tools})),
             __config__=ConfigDict(json_schema_extra=
                                   {
@@ -51,24 +46,32 @@ class SQLToolkit(BaseToolkit):
         )
 
     @classmethod
+    @filter_missconfigured_index_tools
     def get_toolkit(cls, selected_tools: list[str] | None = None, toolkit_name: Optional[str] = None, **kwargs):
         if selected_tools is None:
             selected_tools = []
-        sql_api_wrapper = SQLApiWrapper(**kwargs)
-        prefix = clean_string(toolkit_name, cls.toolkit_max_length) + TOOLKIT_SPLITTER if toolkit_name else ''
+        wrapper_payload = {
+            **kwargs,
+            **kwargs.get('sql_configuration', {}),
+        }
+        sql_api_wrapper = SQLApiWrapper(**wrapper_payload)
         available_tools = sql_api_wrapper.get_available_tools()
         tools = []
         for tool in available_tools:
             if selected_tools and tool["name"] not in selected_tools:
                 continue
+            description = f"{tool['description']}\nDatabase: {sql_api_wrapper.database_name}. Host: {sql_api_wrapper.host}"
+            if toolkit_name:
+                description = f"{description}\nToolkit: {toolkit_name}"
+            description = description[:1000]
             tools.append(BaseAction(
                 api_wrapper=sql_api_wrapper,
-                name=prefix + tool["name"],
-                description=f"{tool['description']}\nDatabase: {sql_api_wrapper.database_name}. Host: {sql_api_wrapper.host}",
-                args_schema=tool["args_schema"]
+                name=tool["name"],
+                description=description,
+                args_schema=tool["args_schema"],
+                metadata={TOOLKIT_NAME_META: toolkit_name, TOOLKIT_TYPE_META: name, TOOL_NAME_META: tool["name"]} if toolkit_name else {TOOL_NAME_META: tool["name"]}
             ))
         return cls(tools=tools)
 
     def get_tools(self) -> list[BaseTool]:
         return self.tools
-

alita_sdk/tools/sql/api_wrapper.py

@@ -1,8 +1,9 @@
 import logging
-from typing import Optional, Any
+from typing import Optional
 
-from pydantic import BaseModel, create_model, model_validator, Field, SecretStr
-from pydantic.fields import PrivateAttr
+from langchain_core.tools import ToolException
+from pydantic import create_model, SecretStr, model_validator
+from pydantic.fields import PrivateAttr, Field
 from sqlalchemy import create_engine, text, inspect, Engine
 from sqlalchemy.orm import sessionmaker
 
@@ -27,7 +28,7 @@ class SQLApiWrapper(BaseToolApiWrapper):
     username: str
     password: SecretStr
     database_name: str
-    _client: Optional[Engine] = PrivateAttr()
+    _client: Optional[Engine] = PrivateAttr(default=None)
 
     @model_validator(mode='before')
     @classmethod
@@ -35,27 +36,73 @@ class SQLApiWrapper(BaseToolApiWrapper):
         for field in SQLConfig.model_fields:
             if field not in values or not values[field]:
                 raise ValueError(f"{field} is a required field and must be provided.")
-
-        dialect = values['dialect']
-        host = values['host']
-        username = values['username']
-        password = values['password']
-        database_name = values['database_name']
-        port = values['port']
-
-        if dialect == SQLDialect.POSTGRES:
-            connection_string = f'postgresql+psycopg2://{username}:{password}@{host}:{port}/{database_name}'
-        elif dialect == SQLDialect.MYSQL:
-            connection_string = f'mysql+pymysql://{username}:{password}@{host}:{port}/{database_name}'
-        else:
-            raise ValueError(f"Unsupported database type. Supported types are: {[e.value for e in SQLDialect]}")
-
-        cls._client = create_engine(connection_string)
         return values
 
+    def _mask_password_in_error(self, error_message: str) -> str:
+        """Mask password in error messages, showing only last 4 characters."""
+        password_str = self.password.get_secret_value()
+        if len(password_str) <= 4:
+            masked_password = "****"
+        else:
+            masked_password = "****" + password_str[-4:]
+
+        # Replace all occurrences of the password, and any substring of the password that may appear in the error message
+        for part in [password_str, password_str.replace('@', ''), password_str.split('@')[-1]]:
+            if part and part in error_message:
+                error_message = error_message.replace(part, masked_password)
+        return error_message
+
+    @property
+    def client(self) -> Engine:
+        """Lazy property to create and return database engine with error handling."""
+        if self._client is None:
+            try:
+                dialect = self.dialect
+                host = self.host
+                username = self.username
+                password = self.password.get_secret_value()
+                database_name = self.database_name
+                port = self.port
+
+                if dialect == SQLDialect.POSTGRES:
+                    connection_string = f'postgresql+psycopg2://{username}:{password}@{host}:{port}/{database_name}'
+                elif dialect == SQLDialect.MYSQL:
+                    connection_string = f'mysql+pymysql://{username}:{password}@{host}:{port}/{database_name}'
+                else:
+                    raise ValueError(f"Unsupported database type. Supported types are: {[e.value for e in SQLDialect]}")
+
+                self._client = create_engine(connection_string)
+
+                # Test the connection
+                with self._client.connect() as conn:
+                    conn.execute(text("SELECT 1"))
+
+            except Exception as e:
+                error_message = str(e)
+                masked_error = self._mask_password_in_error(error_message)
+                logger.error(f"Database connection failed: {masked_error}")
+                raise ValueError(f"Database connection failed: {masked_error}")
+
+        return self._client
+
+    def _handle_database_errors(func):
+        """Decorator to catch exceptions and mask passwords in error messages."""
+
+        def wrapper(self, *args, **kwargs):
+            try:
+                return func(self, *args, **kwargs)
+            except Exception as e:
+                error_message = str(e)
+                masked_error = self._mask_password_in_error(error_message)
+                logger.error(f"Database operation failed in {func.__name__}: {masked_error}")
+                raise ToolException(masked_error)
+
+        return wrapper
+
+    @_handle_database_errors
     def execute_sql(self, sql_query: str):
         """Executes the provided SQL query on the configured database."""
-        engine = self._client
+        engine = self.client
         maker_session = sessionmaker(bind=engine)
         session = maker_session()
         try:
@@ -76,9 +123,10 @@ class SQLApiWrapper(BaseToolApiWrapper):
         finally:
             session.close()
 
+    @_handle_database_errors
     def list_tables_and_columns(self):
         """Lists all tables and their columns in the configured database."""
-        inspector = inspect(self._client)
+        inspector = inspect(self.client)
         data = {}
         tables = inspector.get_table_names()
         for table in tables:
@@ -109,4 +157,4 @@ class SQLApiWrapper(BaseToolApiWrapper):
                 "description": self.list_tables_and_columns.__doc__,
                 "args_schema": SQLNoInput,
             }
-        ]
+        ]

alita_sdk/tools/testio/__init__.py

@@ -1,25 +1,25 @@
 from typing import List, Literal, Optional
 
 from langchain_core.tools import BaseToolkit, BaseTool
-from pydantic import create_model, BaseModel, ConfigDict, Field, SecretStr
+from pydantic import create_model, BaseModel, ConfigDict, Field
 
 from .api_wrapper import TestIOApiWrapper
 from ..base.tool import BaseAction
-from ..utils import clean_string, TOOLKIT_SPLITTER, get_max_toolkit_length
+from ..elitea_base import filter_missconfigured_index_tools
+from ..utils import clean_string, get_max_toolkit_length
+from ...configurations.testio import TestIOConfiguration
+from ...runtime.utils.constants import TOOLKIT_NAME_META, TOOL_NAME_META, TOOLKIT_TYPE_META
 
 name = "testio"
 
 def get_tools(tool):
     return TestIOToolkit().get_toolkit(
         selected_tools=tool['settings'].get('selected_tools', []),
-        endpoint=tool['settings']['endpoint'],
-        api_key=tool['settings']['api_key'],
+        testio_configuration=tool['settings']['testio_configuration'],
         toolkit_name=tool['toolkit_name']
     ).get_tools()
 
 
-TOOLKIT_MAX_LENGTH = 25
-
 class TestIOToolkit(BaseToolkit):
     tools: list[BaseTool] = []
     
@@ -29,8 +29,7 @@ class TestIOToolkit(BaseToolkit):
         selected_tools = {x['name']: x['args_schema'].schema() for x in TestIOApiWrapper.model_construct().get_available_tools()}
         return create_model(
             name,
-            endpoint=(str, Field(description="TestIO endpoint", json_schema_extra={'toolkit_name': True, 'max_toolkit_length': TOOLKIT_MAX_LENGTH})),
-            api_key=(SecretStr, Field(description="API key", json_schema_extra={'secret': True})),
+            testio_configuration=(TestIOConfiguration, Field(description="TestIO Configuration", json_schema_extra={'configuration_types': ['testio']})),
             selected_tools=(List[Literal[tuple(selected_tools)]], Field(default=[], json_schema_extra={'args_schemas': selected_tools})),
             __config__=ConfigDict(json_schema_extra={'metadata': {"label": "TestIO", "icon_url": "testio-icon.svg",
                                                                   "categories": ["testing"],
@@ -38,21 +37,30 @@ class TestIOToolkit(BaseToolkit):
         )
 
     @classmethod
+    @filter_missconfigured_index_tools
     def get_toolkit(cls, selected_tools: list[str] | None = None, toolkit_name: Optional[str] = None, **kwargs):
         if selected_tools is None:
             selected_tools = []
-        testio_api_wrapper = TestIOApiWrapper(**kwargs)
-        prefix = clean_string(toolkit_name, TOOLKIT_MAX_LENGTH) + TOOLKIT_SPLITTER if toolkit_name else ''
+        wrapper_payload = {
+            **kwargs,
+            **kwargs.get('testio_configuration', {}),
+        }
+        testio_api_wrapper = TestIOApiWrapper(**wrapper_payload)
         available_tools = testio_api_wrapper.get_available_tools()
         tools = []
         for tool in available_tools:
             if selected_tools and tool["name"] not in selected_tools:
                 continue
+            description = tool["description"]
+            if toolkit_name:
+                description = f"Toolkit: {toolkit_name}\n{description}"
+            description = description[:1000]
             tools.append(BaseAction(
                 api_wrapper=testio_api_wrapper,
-                name=prefix + tool["name"],
-                description=tool["description"],
-                args_schema=tool["args_schema"]
+                name=tool["name"],
+                description=description,
+                args_schema=tool["args_schema"],
+                metadata={TOOLKIT_NAME_META: toolkit_name, TOOLKIT_TYPE_META: name, TOOL_NAME_META: tool["name"]} if toolkit_name else {TOOL_NAME_META: tool["name"]}
             ))
         return cls(tools=tools)
 

alita_sdk/tools/testrail/__init__.py

@@ -6,16 +6,17 @@ import requests
 
 from .api_wrapper import TestrailAPIWrapper
 from ..base.tool import BaseAction
-from ..utils import clean_string, TOOLKIT_SPLITTER, get_max_toolkit_length, check_connection_response
+from ..elitea_base import filter_missconfigured_index_tools
+from ..utils import clean_string, get_max_toolkit_length, check_connection_response
 from ...configurations.testrail import TestRailConfiguration
 from ...configurations.pgvector import PgVectorConfiguration
+from ...runtime.utils.constants import TOOLKIT_NAME_META, TOOL_NAME_META, TOOLKIT_TYPE_META
 
 name = "testrail"
 
 def get_tools(tool):
     return TestrailToolkit().get_toolkit(
         selected_tools=tool['settings'].get('selected_tools', []),
-        url=tool['settings']['url'],
         testrail_configuration=tool['settings']['testrail_configuration'],
         toolkit_name=tool.get('toolkit_name'),
         llm=tool['settings'].get('llm', None),
@@ -31,17 +32,12 @@ def get_tools(tool):
 
 class TestrailToolkit(BaseToolkit):
     tools: List[BaseTool] = []
-    toolkit_max_length: int = 0
 
     @staticmethod
     def toolkit_config_schema() -> BaseModel:
         selected_tools = {x['name']: x['args_schema'].schema() for x in TestrailAPIWrapper.model_construct().get_available_tools()}
-        TestrailToolkit.toolkit_max_length = get_max_toolkit_length(selected_tools)
         m = create_model(
             name,
-            name=(str, Field(description="Toolkit name", json_schema_extra={
-                'toolkit_name': True,
-                "max_length": TestrailToolkit.toolkit_max_length})),
             testrail_configuration=(Optional[TestRailConfiguration], Field(description="TestRail Configuration", json_schema_extra={'configuration_types': ['testrail']})),
             pgvector_configuration=(Optional[PgVectorConfiguration], Field(default = None,
                                                                            description="PgVector Configuration", json_schema_extra={'configuration_types': ['pgvector']})),
@@ -68,6 +64,7 @@ class TestrailToolkit(BaseToolkit):
         return m
 
     @classmethod
+    @filter_missconfigured_index_tools
     def get_toolkit(cls, selected_tools: list[str] | None = None, toolkit_name: Optional[str] = None, **kwargs):
         if selected_tools is None:
             selected_tools = []
@@ -78,18 +75,23 @@ class TestrailToolkit(BaseToolkit):
             **(kwargs.get('pgvector_configuration') or {}),
         }
         testrail_api_wrapper = TestrailAPIWrapper(**wrapper_payload)
-        prefix = clean_string(toolkit_name, cls.toolkit_max_length) + TOOLKIT_SPLITTER if toolkit_name else ''
         available_tools = testrail_api_wrapper.get_available_tools()
         tools = []
         for tool in available_tools:
             if selected_tools:
                 if tool["name"] not in selected_tools:
                     continue
+            description = tool["description"]
+            if toolkit_name:
+                description = f"Toolkit: {toolkit_name}\n{description}"
+            description = description + "\nTestrail instance: " + testrail_api_wrapper.url
+            description = description[:1000]
             tools.append(BaseAction(
                 api_wrapper=testrail_api_wrapper,
-                name=prefix + tool["name"],
-                description=tool["description"] + "\nTestrail instance: " + testrail_api_wrapper.url,
-                args_schema=tool["args_schema"]
+                name=tool["name"],
+                description=description,
+                args_schema=tool["args_schema"],
+                metadata={TOOLKIT_NAME_META: toolkit_name, TOOLKIT_TYPE_META: name, TOOL_NAME_META: tool["name"]} if toolkit_name else {TOOL_NAME_META: tool["name"]}
             ))
         return cls(tools=tools)