alita-sdk 0.3.366__tar.gz → 0.3.367__tar.gz

{alita_sdk-0.3.366 → alita_sdk-0.3.367}/PKG-INFO

@@ -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 → alita_sdk-0.3.367}/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-0.3.366 → alita_sdk-0.3.367}/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 → alita_sdk-0.3.367}/alita_sdk.egg-info/PKG-INFO

@@ -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 → alita_sdk-0.3.367}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "alita_sdk"
-version = "0.3.366"
+version = "0.3.367"
 description = "SDK for building langchain agents using resources from Alita"
 readme = "README.md"
 requires-python = ">=3.10"