alita-sdk 0.3.365__py3-none-any.whl → 0.3.367__py3-none-any.whl

alita_sdk/runtime/langchain/langraph_agent.py

@@ -786,24 +786,27 @@ class LangGraphAgentRunnable(CompiledStateGraph):
         else:
             result = super().invoke(input, config=config, *args, **kwargs)
         try:
-            if self.output_variables and self.output_variables[0] != "messages":
-                # If output_variables are specified, use the value of first one or use the last messages as default
-                output = result.get(self.output_variables[0])
-                if not output:
-                    output = result['messages'][-1].content
-            else:
-                output = result['messages'][-1].content
+            # if self.output_variables and self.output_variables[0] != "messages":
+            #     # If output_variables are specified, use the value of first one or use the last messages as default
+            #     output = result.get(self.output_variables[0])
+            #     if not output:
+            #         output = result['messages'][-1].content
+            # else:
+            #     output = result['messages'][-1].content
+            output = result['messages'][-1].content
         except:
             output = list(result.values())[-1]
-        thread_id = None
         config_state = self.get_state(config)
-        if config_state.next:
-            thread_id = config['configurable']['thread_id']
+        is_execution_finished = not config_state.next
+        if is_execution_finished:
+            thread_id = None
+
+
 
         result_with_state = {
             "output": output,
             "thread_id": thread_id,
-            "execution_finished": not config_state.next
+            "execution_finished": is_execution_finished
         }
 
         # Include all state values in the result

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-0.3.365.dist-info → alita_sdk-0.3.367.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: alita_sdk
-Version: 0.3.365
+Version: 0.3.367
 Summary: SDK for building langchain agents using resources from Alita
 Author-email: Artem Rozumenko <artyom.rozumenko@gmail.com>, Mikalai Biazruchka <mikalai_biazruchka@epam.com>, Roman Mitusov <roman_mitusov@epam.com>, Ivan Krakhmaliuk <lifedj27@gmail.com>, Artem Dubrovskiy <ad13box@gmail.com>
 License-Expression: Apache-2.0

{alita_sdk-0.3.365.dist-info → alita_sdk-0.3.367.dist-info}/RECORD

@@ -44,7 +44,7 @@ alita_sdk/runtime/langchain/assistant.py,sha256=2tH8je9uKegIIIZUuiGU4zqRVg7jyQas
 alita_sdk/runtime/langchain/chat_message_template.py,sha256=kPz8W2BG6IMyITFDA5oeb5BxVRkHEVZhuiGl4MBZKdc,2176
 alita_sdk/runtime/langchain/constants.py,sha256=eHVJ_beJNTf1WJo4yq7KMK64fxsRvs3lKc34QCXSbpk,3319
 alita_sdk/runtime/langchain/indexer.py,sha256=0ENHy5EOhThnAiYFc7QAsaTNp9rr8hDV_hTK8ahbatk,37592
-alita_sdk/runtime/langchain/langraph_agent.py,sha256=Nrhhcr7fIGaMFwps2XV0ZPL72Ek4wHo7Zml-JaE3AvU,48039
+alita_sdk/runtime/langchain/langraph_agent.py,sha256=n7zhgMSB5p2OwaWu4sZGcJR4RpZNvAO9uNLxN2ImgQA,48109
 alita_sdk/runtime/langchain/mixedAgentParser.py,sha256=M256lvtsL3YtYflBCEp-rWKrKtcY1dJIyRGVv7KW9ME,2611
 alita_sdk/runtime/langchain/mixedAgentRenderes.py,sha256=asBtKqm88QhZRILditjYICwFVKF5KfO38hu2O-WrSWE,5964
 alita_sdk/runtime/langchain/store_manager.py,sha256=i8Fl11IXJhrBXq1F1ukEVln57B1IBe-tqSUvfUmBV4A,2218
@@ -322,7 +322,7 @@ alita_sdk/tools/sharepoint/utils.py,sha256=fZ1YzAu5CTjKSZeslowpOPH974902S8vCp1Wu
 alita_sdk/tools/slack/__init__.py,sha256=YiPAoRc6y6uVpfHl0K1Qi-flcyPlWFIMVcVbhicGWXY,3990
 alita_sdk/tools/slack/api_wrapper.py,sha256=5VrV7iSGno8ZcDzEHdGPNhInhtODGPPvAzoZ9W9iQWE,14009
 alita_sdk/tools/sql/__init__.py,sha256=F3eewPEKqVh3gZdNTUvcoFPOgG9Mn11qKoadtCmhQ4Q,3484
-alita_sdk/tools/sql/api_wrapper.py,sha256=Rky0_CX9HWDQ2mClHGAgP3LHjYVX4iymPuilZMtaDlQ,3687
+alita_sdk/tools/sql/api_wrapper.py,sha256=1VkwnMBr7hGn-LsXXhG774QQQTHvUZilnDuflN2pQes,5810
 alita_sdk/tools/sql/models.py,sha256=AKJgSl_kEEz4fZfw3kbvdGHXaRZ-yiaqfJOB6YOj3i0,641
 alita_sdk/tools/testio/__init__.py,sha256=NEvQtzsffqAXryaffVk0GpdcxZQ1AMkfeztnxHwNql4,2798
 alita_sdk/tools/testio/api_wrapper.py,sha256=BvmL5h634BzG6p7ajnQLmj-uoAw1gjWnd4FHHu1h--Q,21638
@@ -352,8 +352,8 @@ alita_sdk/tools/zephyr_scale/api_wrapper.py,sha256=kT0TbmMvuKhDUZc0i7KO18O38JM9S
 alita_sdk/tools/zephyr_squad/__init__.py,sha256=0ne8XLJEQSLOWfzd2HdnqOYmQlUliKHbBED5kW_Vias,2895
 alita_sdk/tools/zephyr_squad/api_wrapper.py,sha256=kmw_xol8YIYFplBLWTqP_VKPRhL_1ItDD0_vXTe_UuI,14906
 alita_sdk/tools/zephyr_squad/zephyr_squad_cloud_client.py,sha256=R371waHsms4sllHCbijKYs90C-9Yu0sSR3N4SUfQOgU,5066
-alita_sdk-0.3.365.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-alita_sdk-0.3.365.dist-info/METADATA,sha256=MEwGE5tSEKyqZhKFzGM0fQcYLTOzHnrLermRmpJItwU,19071
-alita_sdk-0.3.365.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-alita_sdk-0.3.365.dist-info/top_level.txt,sha256=0vJYy5p_jK6AwVb1aqXr7Kgqgk3WDtQ6t5C-XI9zkmg,10
-alita_sdk-0.3.365.dist-info/RECORD,,
+alita_sdk-0.3.367.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+alita_sdk-0.3.367.dist-info/METADATA,sha256=Vl32B-AoI7amFTAWHAv_2TaEthWlRBcuOm5WiG8XaH4,19071
+alita_sdk-0.3.367.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+alita_sdk-0.3.367.dist-info/top_level.txt,sha256=0vJYy5p_jK6AwVb1aqXr7Kgqgk3WDtQ6t5C-XI9zkmg,10
+alita_sdk-0.3.367.dist-info/RECORD,,