solace-agent-mesh 1.0.5__py3-none-any.whl → 1.0.7__py3-none-any.whl

solace_agent_mesh/agent/adk/intelligent_mcp_callbacks.py

@@ -0,0 +1,414 @@
+"""
+Intelligent MCP Callback Functions
+
+This module contains the refactored MCP callback functions that use intelligent
+content processing to save MCP tool responses as appropriately typed artifacts.
+"""
+
+import json
+import uuid
+from datetime import datetime, timezone
+from typing import Any, Dict, TYPE_CHECKING, List, Optional
+from enum import Enum
+from pydantic import BaseModel
+
+from google.adk.tools import ToolContext, BaseTool
+from solace_ai_connector.common.log import log
+
+from .mcp_content_processor import MCPContentProcessor, MCPContentProcessorConfig
+from ...agent.utils.artifact_helpers import (
+    save_artifact_with_metadata,
+    DEFAULT_SCHEMA_MAX_KEYS,
+)
+from ...agent.utils.context_helpers import get_original_session_id
+
+
+if TYPE_CHECKING:
+    from ...agent.sac.component import SamAgentComponent
+
+
+class McpSaveStatus(str, Enum):
+    """Enumeration for the status of an MCP save operation."""
+
+    SUCCESS = "success"
+    PARTIAL_SUCCESS = "partial_success"
+    ERROR = "error"
+
+
+class SavedArtifactInfo(BaseModel):
+    """
+    A Pydantic model to hold the details of a successfully saved artifact.
+    This mirrors the dictionary structure returned by save_artifact_with_metadata.
+    """
+
+    status: str
+    data_filename: str
+    data_version: int
+    metadata_filename: str
+    metadata_version: int
+    message: str
+
+
+class McpSaveResult(BaseModel):
+    """
+    The definitive, type-safe result of an MCP response save operation.
+
+    Attributes:
+        status: The overall status of the save operation.
+        message: A human-readable summary of the outcome.
+        artifacts_saved: A list of successfully created "intelligent" artifacts.
+        fallback_artifact: An optional artifact representing the raw JSON response,
+                           created only if intelligent processing failed.
+    """
+
+    status: McpSaveStatus
+    message: str
+    artifacts_saved: List[SavedArtifactInfo] = []
+    fallback_artifact: Optional[SavedArtifactInfo] = None
+
+
+async def save_mcp_response_as_artifact_intelligent(
+    tool: BaseTool,
+    tool_context: ToolContext,
+    host_component: "SamAgentComponent",
+    mcp_response_dict: Dict[str, Any],
+    original_tool_args: Dict[str, Any],
+) -> McpSaveResult:
+    """
+    Intelligently processes and saves MCP tool response content as typed artifacts.
+
+    This function uses intelligent content processing to:
+    - Detect and parse different content types (text, image, audio, resource)
+    - Create appropriately typed artifacts with proper MIME types
+    - Generate enhanced metadata based on content analysis
+    - Fall back to raw JSON saving if intelligent processing fails
+
+    Args:
+        tool: The MCPTool instance that generated the response.
+        tool_context: The ADK ToolContext.
+        host_component: The A2A_ADK_HostComponent instance for accessing config and services.
+        mcp_response_dict: The raw MCP tool response dictionary.
+        original_tool_args: The original arguments passed to the MCP tool.
+
+    Returns:
+        An McpSaveResult object containing the structured result of the operation,
+        including status, a list of successfully saved artifacts, and any
+        fallback artifact.
+    """
+    log_identifier = f"[IntelligentMCPCallback:{tool.name}]"
+    log.debug("%s Starting intelligent MCP response artifact saving...", log_identifier)
+
+    processor_config_dict = host_component.get_config("mcp_intelligent_processing", {})
+    processor_config = MCPContentProcessorConfig.from_dict(processor_config_dict)
+
+    saved_artifacts: List[SavedArtifactInfo] = []
+    failed_artifacts: List[Dict[str, Any]] = []
+    fallback_artifact: Optional[SavedArtifactInfo] = None
+    overall_status = McpSaveStatus.SUCCESS
+
+    try:
+        if not processor_config.enable_intelligent_processing:
+            log.info(
+                "%s Intelligent processing disabled, using raw JSON fallback.",
+                log_identifier,
+            )
+            fallback_dict = await _save_raw_mcp_response_fallback(
+                tool,
+                tool_context,
+                host_component,
+                mcp_response_dict,
+                original_tool_args,
+            )
+            if fallback_dict.get("status") in ["success", "partial_success"]:
+                fallback_artifact = SavedArtifactInfo(**fallback_dict)
+                status = McpSaveStatus.SUCCESS
+            else:
+                status = McpSaveStatus.ERROR
+            return McpSaveResult(
+                status=status,
+                message="Intelligent processing disabled; saved raw JSON as fallback.",
+                fallback_artifact=fallback_artifact,
+            )
+
+        processor = MCPContentProcessor(tool.name, original_tool_args)
+        content_items = processor.process_mcp_response(mcp_response_dict)
+
+        if not content_items:
+            log.warning(
+                "%s No content items found, falling back to raw JSON.", log_identifier
+            )
+            fallback_dict = await _save_raw_mcp_response_fallback(
+                tool,
+                tool_context,
+                host_component,
+                mcp_response_dict,
+                original_tool_args,
+            )
+            if fallback_dict.get("status") in ["success", "partial_success"]:
+                fallback_artifact = SavedArtifactInfo(**fallback_dict)
+            return McpSaveResult(
+                status=McpSaveStatus.PARTIAL_SUCCESS,
+                message="No content items found in MCP response; saved raw JSON as fallback.",
+                fallback_artifact=fallback_artifact,
+            )
+
+        log.info(
+            "%s Processing %d content items intelligently.",
+            log_identifier,
+            len(content_items),
+        )
+
+        for item in content_items:
+            try:
+                if hasattr(item, "uri"):
+                    item.uri = str(item.uri)
+                result_dict = await _save_content_item_as_artifact(
+                    item, tool_context, host_component
+                )
+                if result_dict.get("status") in ["success", "partial_success"]:
+                    saved_artifacts.append(SavedArtifactInfo(**result_dict))
+                else:
+                    log.warning(
+                        "%s Failed to save content item: %s",
+                        log_identifier,
+                        result_dict.get("message", "Unknown error"),
+                    )
+                    overall_status = McpSaveStatus.PARTIAL_SUCCESS
+                    failed_artifacts.append(result_dict)
+            except Exception as e:
+                if not processor_config.fallback_to_raw_on_error:
+                    raise
+                log.exception("%s Error saving content item: %s", log_identifier, e)
+                overall_status = McpSaveStatus.PARTIAL_SUCCESS
+                failed_artifacts.append({"status": "error", "message": str(e)})
+
+        if not saved_artifacts:
+            if failed_artifacts:
+                first_error_msg = failed_artifacts[0].get("message", "Unknown error")
+                log.warning(
+                    "%s No items saved successfully. First error: %s",
+                    log_identifier,
+                    first_error_msg,
+                )
+                return McpSaveResult(
+                    status=McpSaveStatus.ERROR,
+                    message=f"Content processing failed. First error: {first_error_msg}",
+                )
+
+            fallback_dict = await _save_raw_mcp_response_fallback(
+                tool,
+                tool_context,
+                host_component,
+                mcp_response_dict,
+                original_tool_args,
+            )
+            if fallback_dict.get("status") in ["success", "partial_success"]:
+                fallback_artifact = SavedArtifactInfo(**fallback_dict)
+            return McpSaveResult(
+                status=McpSaveStatus.PARTIAL_SUCCESS,
+                message="Content processing failed for all items; saved raw JSON as fallback.",
+                fallback_artifact=fallback_artifact,
+            )
+
+        if processor_config_dict.get("save_raw_alongside_intelligent", False):
+            try:
+                fallback_dict = await _save_raw_mcp_response_fallback(
+                    tool,
+                    tool_context,
+                    host_component,
+                    mcp_response_dict,
+                    original_tool_args,
+                )
+                if fallback_dict.get("status") in ["success", "partial_success"]:
+                    fallback_artifact = SavedArtifactInfo(**fallback_dict)
+            except Exception as e:
+                log.warning(
+                    "%s Failed to save raw JSON alongside: %s", log_identifier, e
+                )
+
+        log.info(
+            "%s Intelligent processing complete: %d artifacts saved, status: %s",
+            log_identifier,
+            len(saved_artifacts),
+            overall_status.value,
+        )
+        return McpSaveResult(
+            status=overall_status,
+            artifacts_saved=saved_artifacts,
+            fallback_artifact=fallback_artifact,
+            message=f"Successfully processed {len(saved_artifacts)} content items.",
+        )
+
+    except Exception as e:
+        log.exception(
+            "%s Error in intelligent MCP response processing: %s", log_identifier, e
+        )
+        if processor_config.fallback_to_raw_on_error:
+            log.info(
+                "%s Falling back to raw JSON due to processing error.", log_identifier
+            )
+            try:
+                fallback_dict = await _save_raw_mcp_response_fallback(
+                    tool,
+                    tool_context,
+                    host_component,
+                    mcp_response_dict,
+                    original_tool_args,
+                )
+                if fallback_dict.get("status") in ["success", "partial_success"]:
+                    fallback_artifact = SavedArtifactInfo(**fallback_dict)
+                return McpSaveResult(
+                    status=McpSaveStatus.PARTIAL_SUCCESS,
+                    artifacts_saved=saved_artifacts,
+                    fallback_artifact=fallback_artifact,
+                    message=f"Intelligent processing failed, saved raw JSON as fallback: {e}",
+                )
+            except Exception as fallback_error:
+                log.exception(
+                    "%s Fallback also failed: %s", log_identifier, fallback_error
+                )
+
+        return McpSaveResult(
+            status=McpSaveStatus.ERROR,
+            artifacts_saved=saved_artifacts,
+            fallback_artifact=None,
+            message=f"Failed to save MCP response as artifact: {e}",
+        )
+
+
+async def _save_content_item_as_artifact(
+    content_item,
+    tool_context: ToolContext,
+    host_component: "SamAgentComponent",
+) -> Dict[str, Any]:
+    """Save a single processed content item as an artifact."""
+
+    log_identifier = f"[IntelligentMCPCallback:SaveContentItem:{content_item.filename}]"
+
+    try:
+        artifact_service = host_component.artifact_service
+        if not artifact_service:
+            raise ValueError("ArtifactService is not available on host_component.")
+
+        app_name = host_component.agent_name
+        user_id = tool_context._invocation_context.user_id
+        session_id = get_original_session_id(tool_context._invocation_context)
+        schema_max_keys = host_component.get_config(
+            "schema_max_keys", DEFAULT_SCHEMA_MAX_KEYS
+        )
+        artifact_timestamp = datetime.now(timezone.utc)
+
+        log.debug(
+            "%s Saving content item: type=%s, mime_type=%s, size=%d bytes",
+            log_identifier,
+            content_item.content_type,
+            content_item.mime_type,
+            len(content_item.content_bytes),
+        )
+
+        save_result = await save_artifact_with_metadata(
+            artifact_service=artifact_service,
+            app_name=app_name,
+            user_id=user_id,
+            session_id=session_id,
+            filename=content_item.filename,
+            content_bytes=content_item.content_bytes,
+            mime_type=content_item.mime_type,
+            metadata_dict=content_item.metadata,
+            timestamp=artifact_timestamp,
+            schema_max_keys=schema_max_keys,
+            tool_context=tool_context,
+        )
+
+        log.info(
+            "%s Content item saved as artifact '%s' (version %s). Status: %s",
+            log_identifier,
+            save_result.get("data_filename", content_item.filename),
+            save_result.get("data_version", "N/A"),
+            save_result.get("status"),
+        )
+
+        return save_result
+
+    except Exception as e:
+        log.exception("%s Error saving content item as artifact: %s", log_identifier, e)
+        return {
+            "status": "error",
+            "data_filename": content_item.filename,
+            "message": f"Failed to save content item as artifact: {e}",
+        }
+
+
+async def _save_raw_mcp_response_fallback(
+    tool: BaseTool,
+    tool_context: ToolContext,
+    host_component: "SamAgentComponent",
+    mcp_response_dict: Dict[str, Any],
+    original_tool_args: Dict[str, Any],
+) -> Dict[str, Any]:
+    """
+    Fallback function to save the raw MCP response as a JSON artifact.
+    This is the original behavior, used when intelligent processing is disabled or fails.
+    """
+    log_identifier = f"[IntelligentMCPCallback:{tool.name}:RawFallback]"
+    log.debug("%s Saving raw MCP response as JSON artifact...", log_identifier)
+
+    try:
+        a2a_context = tool_context.state.get("a2a_context", {})
+        logical_task_id = a2a_context.get("logical_task_id", "unknownTask")
+        task_id_suffix = logical_task_id[-6:]
+        random_suffix = uuid.uuid4().hex[:6]
+        filename = f"{task_id_suffix}_{tool.name}_raw_{random_suffix}.json"
+
+        content_bytes = json.dumps(mcp_response_dict, indent=2).encode("utf-8")
+        mime_type = "application/json"
+        artifact_timestamp = datetime.now(timezone.utc)
+
+        metadata_for_saving = {
+            "description": f"Raw JSON response from MCP tool {tool.name}",
+            "source_tool_name": tool.name,
+            "source_tool_args": original_tool_args,
+            "processing_type": "raw_fallback",
+        }
+
+        artifact_service = host_component.artifact_service
+        app_name = host_component.agent_name
+        user_id = tool_context._invocation_context.user_id
+        session_id = get_original_session_id(tool_context._invocation_context)
+        schema_max_keys = host_component.get_config(
+            "schema_max_keys", DEFAULT_SCHEMA_MAX_KEYS
+        )
+
+        save_result = await save_artifact_with_metadata(
+            artifact_service=artifact_service,
+            app_name=app_name,
+            user_id=user_id,
+            session_id=session_id,
+            filename=filename,
+            content_bytes=content_bytes,
+            mime_type=mime_type,
+            metadata_dict=metadata_for_saving,
+            timestamp=artifact_timestamp,
+            schema_max_keys=schema_max_keys,
+            tool_context=tool_context,
+        )
+
+        log.info(
+            "%s Raw MCP response saved as artifact '%s' (version %s). Status: %s",
+            log_identifier,
+            save_result.get("data_filename", filename),
+            save_result.get("data_version", "N/A"),
+            save_result.get("status"),
+        )
+
+        return save_result
+
+    except Exception as e:
+        log.exception(
+            "%s Error saving raw MCP response as artifact: %s", log_identifier, e
+        )
+        return {
+            "status": "error",
+            "data_filename": filename if "filename" in locals() else "unknown_filename",
+            "message": f"Failed to save raw MCP response as artifact: {e}",
+        }