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

massgen/backend/chat_completions.py

@@ -31,7 +31,11 @@ from ..logger_config import log_backend_agent_message, log_stream_chunk, logger
 
 # Local imports
 from .base import FilesystemSupport, StreamChunk
-from .base_with_custom_tool_and_mcp import CustomToolAndMCPBackend
+from .base_with_custom_tool_and_mcp import (
+    CustomToolAndMCPBackend,
+    CustomToolChunk,
+    ToolExecutionConfig,
+)
 
 
 class ChatCompletionsBackend(CustomToolAndMCPBackend):
@@ -67,6 +71,71 @@ class ChatCompletionsBackend(CustomToolAndMCPBackend):
         async for chunk in super().stream_with_tools(messages, tools, **kwargs):
             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 Chat Completions 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"
+        """
+        function_output_msg = {
+            "role": "tool",
+            "tool_call_id": call.get("call_id", ""),
+            "content": 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 Chat Completions 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"
+        """
+        error_output_msg = {
+            "role": "tool",
+            "tool_call_id": call.get("call_id", ""),
+            "content": 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]],
@@ -193,18 +262,8 @@ class ChatCompletionsBackend(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 base helper
+            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:
@@ -221,8 +280,8 @@ class ChatCompletionsBackend(CustomToolAndMCPBackend):
                     content="⚠️ [MCP] All servers blocked by circuit breaker",
                     source="circuit_breaker",
                 )
-                yield StreamChunk(type="done")
-                return
+                # Skip MCP tool execution but continue with custom tools
+                mcp_calls = []
 
             # Initialize for execution
             functions_executed = False
@@ -243,8 +302,7 @@ class ChatCompletionsBackend(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"[ChatCompletions] Planning mode enabled - selective blocking of {len(blocked_tools)} tools")
@@ -274,188 +332,44 @@ class ChatCompletionsBackend(CustomToolAndMCPBackend):
                     }
                     updated_messages.append(assistant_message)
 
-            # Execute custom tools first
-            for call in custom_calls:
-                try:
-                    # Yield custom tool call status
-                    yield StreamChunk(
-                        type="custom_tool_status",
-                        status="custom_tool_called",
-                        content=f"🔧 [Custom Tool] Calling {call['name']}...",
-                        source=f"custom_{call['name']}",
-                    )
-
-                    # Yield custom tool arguments
-                    yield StreamChunk(
-                        type="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 result to messages
-                    function_output_msg = {
-                        "role": "tool",
-                        "tool_call_id": call["call_id"],
-                        "content": str(result),
-                    }
-                    updated_messages.append(function_output_msg)
-
-                    # Yield custom tool results
-                    yield StreamChunk(
-                        type="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 StreamChunk(
-                        type="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']}")
-
-                except Exception as e:
-                    logger.error(f"Error executing custom tool {call['name']}: {e}")
-                    error_msg = f"Error executing {call['name']}: {str(e)}"
-
-                    # Yield error with arguments shown
-                    yield StreamChunk(
-                        type="custom_tool_status",
-                        status="function_call",
-                        content=f"Arguments for Calling {call['name']}: {call['arguments']}",
-                        source=f"custom_{call['name']}",
-                    )
+            # Create tool execution configuration objects
+            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,
+            )
 
-                    yield StreamChunk(
-                        type="custom_tool_status",
-                        status="custom_tool_error",
-                        content=f"❌ [Custom Tool Error] {error_msg}",
-                        source=f"custom_{call['name']}",
-                    )
+            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,
+            )
 
-                    # Add error result to messages
-                    error_output_msg = {
-                        "role": "tool",
-                        "tool_call_id": call["call_id"],
-                        "content": error_msg,
-                    }
-                    updated_messages.append(error_output_msg)
-                    processed_call_ids.add(call["call_id"])
-                    functions_executed = True
+            # 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 chunk
+                functions_executed = True
 
-            # Execute MCP function calls
-            mcp_functions_executed = False
+            # Execute MCP tools using unified method
             for call in mcp_calls:
-                function_name = call["name"]
-                if function_name in self._mcp_functions:
-                    yield StreamChunk(
-                        type="mcp_status",
-                        status="mcp_tool_called",
-                        content=f"🔧 [MCP Tool] Calling {function_name}...",
-                        source=f"mcp_{function_name}",
-                    )
-
-                    # Yield detailed MCP status as StreamChunk
-                    tools_info = f" ({len(self._mcp_functions)} tools available)" if self._mcp_functions else ""
-                    yield StreamChunk(
-                        type="mcp_status",
-                        status="mcp_tools_initiated",
-                        content=f"MCP tool call initiated (call #{self._mcp_tool_calls_count}){tools_info}: {function_name}",
-                        source=f"mcp_{function_name}",
-                    )
-
-                    try:
-                        # Execute MCP function with retry and exponential backoff
-                        result_str, result_obj = await self._execute_mcp_function_with_retry(function_name, call["arguments"])
-
-                        # Check if function failed after all retries
-                        if isinstance(result_str, str) and result_str.startswith("Error:"):
-                            # Log failure but still create tool response
-                            logger.warning(f"MCP function {function_name} failed after retries: {result_str}")
-
-                            # Add error result to messages
-                            function_output_msg = {
-                                "role": "tool",
-                                "tool_call_id": call["call_id"],
-                                "content": result_str,
-                            }
-                            updated_messages.append(function_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_output_msg = {
-                            "role": "tool",
-                            "tool_call_id": call["call_id"],
-                            "content": error_msg,
-                        }
-                        updated_messages.append(function_output_msg)
-
-                        processed_call_ids.add(call["call_id"])
-                        mcp_functions_executed = True
-                        continue
-
-                    # Yield function_call status
-                    yield StreamChunk(
-                        type="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 = {
-                        "role": "tool",
-                        "tool_call_id": call["call_id"],
-                        "content": str(result_str),
-                    }
-                    updated_messages.append(function_output_msg)
-
-                    # Yield function_call_output status with preview
-                    result_text = str(result_str)
-                    if hasattr(result_obj, "content") and result_obj.content:
-                        if isinstance(result_obj.content, list) and len(result_obj.content) > 0:
-                            first_item = result_obj.content[0]
-                            if hasattr(first_item, "text"):
-                                result_text = first_item.text
-
-                    yield StreamChunk(
-                        type="mcp_status",
-                        status="function_call_output",
-                        content=f"Results for Calling {function_name}: {result_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 StreamChunk(
-                        type="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
+                async for chunk in self._execute_tool_with_logging(call, mcp_tool_config, updated_messages, processed_call_ids):
+                    yield chunk
+                functions_executed = True
 
             # Ensure all captured function calls have results to prevent hanging
             for call in captured_function_calls:
@@ -469,10 +383,10 @@ class ChatCompletionsBackend(CustomToolAndMCPBackend):
                         "content": 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 = self._trim_message_history(updated_messages)
 
                 # Recursive call with updated messages