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

alita_sdk/runtime/langchain/langraph_agent.py

@@ -786,13 +786,14 @@ 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]
         config_state = self.get_state(config)

alita_sdk/tools/sql/api_wrapper.py

@@ -1,6 +1,7 @@
 import logging
 from typing import Optional
 
+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
@@ -45,8 +46,11 @@ class SQLApiWrapper(BaseToolApiWrapper):
         else:
             masked_password = "****" + password_str[-4:]
 
-        # Replace password in error message
-        return error_message.replace(password_str, masked_password)
+        # 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:
@@ -81,6 +85,21 @@ class SQLApiWrapper(BaseToolApiWrapper):
 
         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
@@ -104,6 +123,7 @@ 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)
@@ -137,4 +157,4 @@ class SQLApiWrapper(BaseToolApiWrapper):
                 "description": self.list_tables_and_columns.__doc__,
                 "args_schema": SQLNoInput,
             }
-        ]
+        ]

{alita_sdk-0.3.366.dist-info → alita_sdk-0.3.367.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: alita_sdk
-Version: 0.3.366
+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.366.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=bQVT70vfIpfAjuqFbVx670XBIdlcvwsXUSltWh2vvxk,48043
+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=AHv90ppE2AYB0z8KXeUNUtnC0Bm-g8vtsLMY9GdXP84,4893
+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.366.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-alita_sdk-0.3.366.dist-info/METADATA,sha256=wlps-YB4WTVcMnF8BT-mKcSZPxM8bMLQUpPkG1kgaeY,19071
-alita_sdk-0.3.366.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-alita_sdk-0.3.366.dist-info/top_level.txt,sha256=0vJYy5p_jK6AwVb1aqXr7Kgqgk3WDtQ6t5C-XI9zkmg,10
-alita_sdk-0.3.366.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,,