aixtools 0.2.16__tar.gz → 0.2.18__tar.gz

{aixtools-0.2.16 → aixtools-0.2.18}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aixtools
-Version: 0.2.16
+Version: 0.2.18
 Summary: Tools for AI exploration and debugging
 Requires-Python: >=3.11.2
 Description-Content-Type: text/markdown

{aixtools-0.2.16 → aixtools-0.2.18}/aixtools/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.2.16'
-__version_tuple__ = version_tuple = (0, 2, 16)
+__version__ = version = '0.2.18'
+__version_tuple__ = version_tuple = (0, 2, 18)
 
-__commit_id__ = commit_id = 'gbf6656c3c'
+__commit_id__ = commit_id = 'gda0a23c48'

{aixtools-0.2.16 → aixtools-0.2.18}/aixtools/agents/agent.py

@@ -105,7 +105,7 @@ def get_model(model_family=MODEL_FAMILY, model_name=None, http_client=None, **kw
         case "azure":
             return _get_model_openai_azure(model_name=model_name or AZURE_MODEL_NAME, http_client=http_client, **kwargs)
         case "bedrock":
-            return _get_model_bedrock(model_name=model_name or BEDROCK_MODEL_NAME, http_client=http_client, **kwargs)
+            return _get_model_bedrock(model_name=model_name or BEDROCK_MODEL_NAME, **kwargs)
         case "ollama":
             return _get_model_ollama(model_name=model_name or OLLAMA_MODEL_NAME, http_client=http_client, **kwargs)
         case "openai":

aixtools-0.2.18/aixtools/logging/mcp_middleware.py

@@ -0,0 +1,54 @@
+"""
+Custom middleware for MCP servers
+"""
+
+import traceback
+
+from fastmcp.server.middleware.error_handling import ErrorHandlingMiddleware
+from fastmcp.server.middleware.middleware import MiddlewareContext
+
+from aixtools.mcp import AixToolError
+
+
+class AixErrorHandlingMiddleware(ErrorHandlingMiddleware):
+    """Custom middleware class for handling errors in MCP servers."""
+
+    def log_as_warn_with_traceback(
+        self, *, error: Exception, original_error: Exception, context: MiddlewareContext
+    ) -> None:
+        """Logs provided error as warning.
+        original_error is an 'unwrapped' error if applicable, otherwise can be the same as error param.
+        """
+        error_type = type(original_error).__name__
+        method = context.method or "unknown"
+        error_key = f"{error_type}:{method}"
+        self.error_counts[error_key] = self.error_counts.get(error_key, 0) + 1
+        base_message = f"{method} resulted in {error_type}: {str(error)}"
+        if self.include_traceback:
+            self.logger.warning(f"{base_message}\n{traceback.format_exc()}")
+        if self.error_callback:
+            try:
+                self.error_callback(error, context)
+            except Exception as callback_error:  # pylint: disable=broad-exception-caught
+                self.logger.warning("Callback failed spectacularly: %s", callback_error)
+
+    def handle_error(self, error: Exception, context: MiddlewareContext) -> bool:
+        """Custom error logging"""
+        if isinstance(error, AixToolError):
+            self.log_as_warn_with_traceback(error=error, original_error=error, context=context)
+            return True
+
+        inner_error = error
+        while hasattr(inner_error, "__cause__") and inner_error.__cause__ is not None:
+            inner_error = inner_error.__cause__
+            if isinstance(inner_error, AixToolError):
+                self.log_as_warn_with_traceback(error=error, original_error=inner_error, context=context)
+                return True
+
+        return False
+
+    def _log_error(self, error: Exception, context: MiddlewareContext) -> None:
+        """Override original _log_error method."""
+        if self.handle_error(error, context):
+            return
+        super()._log_error(error, context)

{aixtools-0.2.16 → aixtools-0.2.18}/aixtools.egg-info/SOURCES.txt

@@ -61,6 +61,7 @@ aixtools/logging/log_objects.py
 aixtools/logging/logging_config.py
 aixtools/logging/mcp_log_models.py
 aixtools/logging/mcp_logger.py
+aixtools/logging/mcp_middleware.py
 aixtools/logging/model_patch_logging.py
 aixtools/logging/open_telemetry.py
 aixtools/mcp/__init__.py