massgen 0.1.4__py3-none-any.whl → 0.1.6__py3-none-any.whl

massgen/backend/grok.py

@@ -20,11 +20,12 @@ from __future__ import annotations
 
 import logging
 import os
-from typing import Any, Dict, List, Optional
+from typing import Any, AsyncGenerator, Dict, List, Optional
 
 from openai import AsyncOpenAI
 
 from ..logger_config import log_stream_chunk
+from .base import StreamChunk
 from .chat_completions import ChatCompletionsBackend
 
 logger = logging.getLogger(__name__)
@@ -70,7 +71,7 @@ class GrokBackend(ChatCompletionsBackend):
         tools: List[Dict[str, Any]],
         client,
         **kwargs,
-    ):
+    ) -> AsyncGenerator[StreamChunk, None]:
         """Override to add Grok-specific search parameters before API call."""
         # Build API params using parent method
         all_params = {**self.config, **kwargs}

massgen/backend/response.py

@@ -23,7 +23,12 @@ from ..formatter import ResponseFormatter
 from ..logger_config import log_backend_agent_message, log_stream_chunk, logger
 from ..stream_chunk import ChunkType, TextStreamChunk
 from .base import FilesystemSupport, StreamChunk
-from .base_with_custom_tool_and_mcp import CustomToolAndMCPBackend, UploadFileError
+from .base_with_custom_tool_and_mcp import (
+    CustomToolAndMCPBackend,
+    CustomToolChunk,
+    ToolExecutionConfig,
+    UploadFileError,
+)
 
 
 class ResponseBackend(CustomToolAndMCPBackend):
@@ -138,6 +143,97 @@ class ResponseBackend(CustomToolAndMCPBackend):
         async for chunk in self._process_stream(stream, all_params, agent_id):
             yield chunk
 
+    def _append_tool_result_message(
+        self,
+        updated_messages: List[Dict[str, Any]],
+        call: Dict[str, Any],
+        result: Any,
+        tool_type: str,
+    ) -> None:
+        """Append tool result to messages in Response API format.
+
+        Args:
+            updated_messages: Message list to append to
+            call: Tool call dictionary with call_id, name, arguments
+            result: Tool execution result
+            tool_type: "custom" or "mcp"
+
+        Note:
+            Response API requires both function_call and function_call_output messages.
+        """
+        # Add function call message
+        function_call_msg = {
+            "type": "function_call",
+            "call_id": call.get("call_id", ""),
+            "name": call.get("name", ""),
+            "arguments": call.get("arguments", "{}"),
+        }
+        updated_messages.append(function_call_msg)
+
+        # Add function output message
+        function_output_msg = {
+            "type": "function_call_output",
+            "call_id": call.get("call_id", ""),
+            "output": str(result),
+        }
+        updated_messages.append(function_output_msg)
+
+    def _append_tool_error_message(
+        self,
+        updated_messages: List[Dict[str, Any]],
+        call: Dict[str, Any],
+        error_msg: str,
+        tool_type: str,
+    ) -> None:
+        """Append tool error to messages in Response API format.
+
+        Args:
+            updated_messages: Message list to append to
+            call: Tool call dictionary with call_id, name, arguments
+            error_msg: Error message string
+            tool_type: "custom" or "mcp"
+
+        Note:
+            Response API requires both function_call and function_call_output messages.
+        """
+        # Add function call message
+        function_call_msg = {
+            "type": "function_call",
+            "call_id": call.get("call_id", ""),
+            "name": call.get("name", ""),
+            "arguments": call.get("arguments", "{}"),
+        }
+        updated_messages.append(function_call_msg)
+
+        # Add error output message
+        error_output_msg = {
+            "type": "function_call_output",
+            "call_id": call.get("call_id", ""),
+            "output": error_msg,
+        }
+        updated_messages.append(error_output_msg)
+
+    async def _execute_custom_tool(self, call: Dict[str, Any]) -> AsyncGenerator[CustomToolChunk, None]:
+        """Execute custom tool with streaming support - async generator for base class.
+
+        This method is called by _execute_tool_with_logging and yields CustomToolChunk
+        objects for intermediate streaming output. The base class detects the async
+        generator and streams intermediate results to users in real-time.
+
+        Args:
+            call: Tool call dictionary with name and arguments
+
+        Yields:
+            CustomToolChunk objects with streaming data
+
+        Note:
+            - Intermediate chunks (completed=False) are streamed to users in real-time
+            - Final chunk (completed=True) contains the accumulated result for message history
+            - The base class automatically handles extracting and displaying intermediate chunks
+        """
+        async for chunk in self.stream_custom_tool_execution(call):
+            yield chunk
+
     async def _stream_with_custom_and_mcp_tools(
         self,
         current_messages: List[Dict[str, Any]],
@@ -224,18 +320,8 @@ class ResponseBackend(CustomToolAndMCPBackend):
 
         # Execute any captured function calls
         if captured_function_calls and response_completed:
-            # Categorize function calls
-            mcp_calls = []
-            custom_calls = []
-            provider_calls = []
-
-            for call in captured_function_calls:
-                if call["name"] in self._mcp_functions:
-                    mcp_calls.append(call)
-                elif call["name"] in self._custom_tool_names:
-                    custom_calls.append(call)
-                else:
-                    provider_calls.append(call)
+            # Categorize function calls using helper method
+            mcp_calls, custom_calls, provider_calls = self._categorize_tool_calls(captured_function_calls)
 
             # If there are provider calls (non-MCP, non-custom), let API handle them
             if provider_calls:
@@ -248,100 +334,54 @@ class ResponseBackend(CustomToolAndMCPBackend):
             updated_messages = current_messages.copy()
             processed_call_ids = set()  # Initialize processed_call_ids here
 
-            # Execute custom tools first
-            for call in custom_calls:
-                try:
-                    # Yield custom tool call status
-                    yield TextStreamChunk(
-                        type=ChunkType.CUSTOM_TOOL_STATUS,
-                        status="custom_tool_called",
-                        content=f"🔧 [Custom Tool] Calling {call['name']}...",
-                        source=f"custom_{call['name']}",
-                    )
-
-                    # Yield custom tool arguments (like MCP tools)
-                    yield TextStreamChunk(
-                        type=ChunkType.CUSTOM_TOOL_STATUS,
-                        status="function_call",
-                        content=f"Arguments for Calling {call['name']}: {call['arguments']}",
-                        source=f"custom_{call['name']}",
-                    )
-
-                    # Execute custom tool
-                    result = await self._execute_custom_tool(call)
-
-                    # Add function call and result to messages
-                    function_call_msg = {
-                        "type": "function_call",
-                        "call_id": call["call_id"],
-                        "name": call["name"],
-                        "arguments": call["arguments"],
-                    }
-                    updated_messages.append(function_call_msg)
-
-                    function_output_msg = {
-                        "type": "function_call_output",
-                        "call_id": call["call_id"],
-                        "output": str(result),
-                    }
-                    updated_messages.append(function_output_msg)
-
-                    # Yield custom tool results (like MCP tools)
-                    yield TextStreamChunk(
-                        type=ChunkType.CUSTOM_TOOL_STATUS,
-                        status="function_call_output",
-                        content=f"Results for Calling {call['name']}: {str(result)}",
-                        source=f"custom_{call['name']}",
-                    )
-
-                    # Yield custom tool response status
-                    yield TextStreamChunk(
-                        type=ChunkType.CUSTOM_TOOL_STATUS,
-                        status="custom_tool_response",
-                        content=f"✅ [Custom Tool] {call['name']} completed",
-                        source=f"custom_{call['name']}",
-                    )
-
-                    processed_call_ids.add(call["call_id"])
-                    functions_executed = True
-                    logger.info(f"Executed custom tool: {call['name']}")
+            # Configuration for custom tool execution
+            CUSTOM_TOOL_CONFIG = ToolExecutionConfig(
+                tool_type="custom",
+                chunk_type="custom_tool_status",
+                emoji_prefix="🔧 [Custom Tool]",
+                success_emoji="✅ [Custom Tool]",
+                error_emoji="❌ [Custom Tool Error]",
+                source_prefix="custom_",
+                status_called="custom_tool_called",
+                status_response="custom_tool_response",
+                status_error="custom_tool_error",
+                execution_callback=self._execute_custom_tool,
+            )
 
-                except Exception as e:
-                    logger.error(f"Error executing custom tool {call['name']}: {e}")
-                    error_msg = f"Error executing {call['name']}: {str(e)}"
+            # Configuration for MCP tool execution
+            MCP_TOOL_CONFIG = ToolExecutionConfig(
+                tool_type="mcp",
+                chunk_type="mcp_status",
+                emoji_prefix="🔧 [MCP Tool]",
+                success_emoji="✅ [MCP Tool]",
+                error_emoji="❌ [MCP Tool Error]",
+                source_prefix="mcp_",
+                status_called="mcp_tool_called",
+                status_response="mcp_tool_response",
+                status_error="mcp_tool_error",
+                execution_callback=self._execute_mcp_function_with_retry,
+            )
 
-                    # Yield error with arguments shown
-                    yield TextStreamChunk(
-                        type=ChunkType.CUSTOM_TOOL_STATUS,
-                        status="function_call",
-                        content=f"Arguments for Calling {call['name']}: {call['arguments']}",
-                        source=f"custom_{call['name']}",
-                    )
+            chunk_type_map = {
+                "custom_tool_status": ChunkType.CUSTOM_TOOL_STATUS,
+                "mcp_status": ChunkType.MCP_STATUS,
+            }
 
+            # Execute custom tools using unified method
+            for call in custom_calls:
+                async for chunk in self._execute_tool_with_logging(
+                    call,
+                    CUSTOM_TOOL_CONFIG,
+                    updated_messages,
+                    processed_call_ids,
+                ):
                     yield TextStreamChunk(
-                        type=ChunkType.CUSTOM_TOOL_STATUS,
-                        status="custom_tool_error",
-                        content=f"❌ [Custom Tool Error] {error_msg}",
-                        source=f"custom_{call['name']}",
+                        type=chunk_type_map.get(chunk.type, chunk.type),
+                        status=getattr(chunk, "status", None),
+                        content=getattr(chunk, "content", None),
+                        source=getattr(chunk, "source", None),
                     )
-
-                    # Add error result to messages
-                    function_call_msg = {
-                        "type": "function_call",
-                        "call_id": call["call_id"],
-                        "name": call["name"],
-                        "arguments": call["arguments"],
-                    }
-                    updated_messages.append(function_call_msg)
-
-                    error_output_msg = {
-                        "type": "function_call_output",
-                        "call_id": call["call_id"],
-                        "output": error_msg,
-                    }
-                    updated_messages.append(error_output_msg)
-                    processed_call_ids.add(call["call_id"])
-                    functions_executed = True
+                functions_executed = True
 
             # Check circuit breaker status before executing MCP functions
             if mcp_calls and not await super()._check_circuit_breaker_before_execution():
@@ -352,11 +392,8 @@ class ResponseBackend(CustomToolAndMCPBackend):
                     content="⚠️ [MCP] All servers blocked by circuit breaker",
                     source="circuit_breaker",
                 )
-                yield TextStreamChunk(type=ChunkType.DONE, source="response_api")
-                return
-
-            # Execute MCP function calls
-            mcp_functions_executed = False
+                # Skip MCP tool execution but continue with custom tool results
+                mcp_calls = []
 
             # Check if planning mode is enabled - selectively block MCP tool execution during planning
             if self.is_planning_mode_enabled():
@@ -372,123 +409,26 @@ class ResponseBackend(CustomToolAndMCPBackend):
                         source="planning_mode",
                     )
                     # Skip all MCP tool execution but still continue with workflow
-                    yield StreamChunk(type="done")
-                    return
+                    mcp_calls = []
                 else:
                     # Selective blocking - log but continue to check each tool individually
                     logger.info(f"[Response] Planning mode enabled - selective blocking of {len(blocked_tools)} tools")
 
-            # Ensure every captured function call gets a result to prevent hanging
-            for call in captured_function_calls:
-                function_name = call["name"]
-                if function_name in self._mcp_functions:
-                    # Yield MCP tool call status
+            # Execute MCP tools using unified method
+            for call in mcp_calls:
+                async for chunk in self._execute_tool_with_logging(
+                    call,
+                    MCP_TOOL_CONFIG,
+                    updated_messages,
+                    processed_call_ids,
+                ):
                     yield TextStreamChunk(
-                        type=ChunkType.MCP_STATUS,
-                        status="mcp_tool_called",
-                        content=f"🔧 [MCP Tool] Calling {function_name}...",
-                        source=f"mcp_{function_name}",
+                        type=chunk_type_map.get(chunk.type, chunk.type),
+                        status=getattr(chunk, "status", None),
+                        content=getattr(chunk, "content", None),
+                        source=getattr(chunk, "source", None),
                     )
-
-                    try:
-                        # Execute MCP function with retry and exponential backoff
-                        result, result_obj = await super()._execute_mcp_function_with_retry(
-                            function_name,
-                            call["arguments"],
-                        )
-
-                        # Check if function failed after all retries
-                        if isinstance(result, str) and result.startswith("Error:"):
-                            # Log failure but still create tool response
-                            logger.warning(f"MCP function {function_name} failed after retries: {result}")
-
-                            # Add error result to messages
-                            function_call_msg = {
-                                "type": "function_call",
-                                "call_id": call["call_id"],
-                                "name": function_name,
-                                "arguments": call["arguments"],
-                            }
-                            updated_messages.append(function_call_msg)
-
-                            error_output_msg = {
-                                "type": "function_call_output",
-                                "call_id": call["call_id"],
-                                "output": result,
-                            }
-                            updated_messages.append(error_output_msg)
-
-                            processed_call_ids.add(call["call_id"])
-                            mcp_functions_executed = True
-                            continue
-
-                    except Exception as e:
-                        # Only catch unexpected non-MCP system errors
-                        logger.error(f"Unexpected error in MCP function execution: {e}")
-                        error_msg = f"Error executing {function_name}: {str(e)}"
-
-                        # Add error result to messages
-                        function_call_msg = {
-                            "type": "function_call",
-                            "call_id": call["call_id"],
-                            "name": function_name,
-                            "arguments": call["arguments"],
-                        }
-                        updated_messages.append(function_call_msg)
-
-                        error_output_msg = {
-                            "type": "function_call_output",
-                            "call_id": call["call_id"],
-                            "output": error_msg,
-                        }
-                        updated_messages.append(error_output_msg)
-
-                        processed_call_ids.add(call["call_id"])
-                        mcp_functions_executed = True
-                        continue
-
-                    # Add function call to messages and yield status chunk
-                    function_call_msg = {
-                        "type": "function_call",
-                        "call_id": call["call_id"],
-                        "name": function_name,
-                        "arguments": call["arguments"],
-                    }
-                    updated_messages.append(function_call_msg)
-                    yield TextStreamChunk(
-                        type=ChunkType.MCP_STATUS,
-                        status="function_call",
-                        content=f"Arguments for Calling {function_name}: {call['arguments']}",
-                        source=f"mcp_{function_name}",
-                    )
-
-                    # Add function output to messages and yield status chunk
-                    function_output_msg = {
-                        "type": "function_call_output",
-                        "call_id": call["call_id"],
-                        "output": str(result),
-                    }
-                    updated_messages.append(function_output_msg)
-                    yield TextStreamChunk(
-                        type=ChunkType.MCP_STATUS,
-                        status="function_call_output",
-                        content=f"Results for Calling {function_name}: {str(result_obj.content[0].text)}",
-                        source=f"mcp_{function_name}",
-                    )
-
-                    logger.info(f"Executed MCP function {function_name} (stdio/streamable-http)")
-                    processed_call_ids.add(call["call_id"])
-
-                    # Yield MCP tool response status
-                    yield TextStreamChunk(
-                        type=ChunkType.MCP_STATUS,
-                        status="mcp_tool_response",
-                        content=f"✅ [MCP Tool] {function_name} completed",
-                        source=f"mcp_{function_name}",
-                    )
-
-                    mcp_functions_executed = True
-                    functions_executed = True
+                functions_executed = True
 
             # Ensure all captured function calls have results to prevent hanging
             for call in captured_function_calls:
@@ -510,10 +450,10 @@ class ResponseBackend(CustomToolAndMCPBackend):
                         "output": f"Error: Tool call {call['call_id']} for function {call['name']} was not processed. This may indicate a validation or execution error.",
                     }
                     updated_messages.append(error_output_msg)
-                    mcp_functions_executed = True
+                    functions_executed = True
 
             # Trim history after function executions to bound memory usage
-            if functions_executed or mcp_functions_executed:
+            if functions_executed:
                 updated_messages = super()._trim_message_history(updated_messages)
 
                 # Recursive call with updated messages