massgen 0.1.0a3__py3-none-any.whl → 0.1.1__py3-none-any.whl

massgen/backend/gemini_trackers.py

@@ -0,0 +1,344 @@
+# -*- coding: utf-8 -*-
+"""
+MCP tracking utilities for the Gemini backend, handling deduplication across streaming chunks and extraction from SDK objects.
+"""
+
+import hashlib
+import json
+import time
+from typing import Any, Dict, Optional
+
+
+class MCPResponseTracker:
+    """
+    Tracks MCP tool responses across streaming chunks to handle deduplication.
+
+    Similar to MCPCallTracker but for tracking tool responses to avoid duplicate output.
+    """
+
+    def __init__(self):
+        """Initialize the tracker with empty storage."""
+        self.processed_responses = set()  # Store hashes of processed responses
+        self.response_history = []  # Store all unique responses with timestamps
+
+    def get_response_hash(self, tool_name: str, tool_response: Any) -> str:
+        """
+        Generate a unique hash for a tool response based on name and response content.
+
+        Args:
+            tool_name: Name of the tool that responded
+            tool_response: Response from the tool
+
+        Returns:
+            MD5 hash string identifying this specific response
+        """
+        # Create a deterministic string representation
+        content = f"{tool_name}:{str(tool_response)}"
+        return hashlib.md5(content.encode()).hexdigest()
+
+    def is_new_response(self, tool_name: str, tool_response: Any) -> bool:
+        """
+        Check if this is a new tool response we haven't seen before.
+
+        Args:
+            tool_name: Name of the tool that responded
+            tool_response: Response from the tool
+
+        Returns:
+            True if this is a new response, False if already processed
+        """
+        response_hash = self.get_response_hash(tool_name, tool_response)
+        return response_hash not in self.processed_responses
+
+    def add_response(self, tool_name: str, tool_response: Any) -> Dict[str, Any]:
+        """
+        Add a new response to the tracker.
+
+        Args:
+            tool_name: Name of the tool that responded
+            tool_response: Response from the tool
+
+        Returns:
+            Dictionary containing response details and timestamp
+        """
+        response_hash = self.get_response_hash(tool_name, tool_response)
+        self.processed_responses.add(response_hash)
+
+        record = {
+            "tool_name": tool_name,
+            "response": tool_response,
+            "hash": response_hash,
+            "timestamp": time.time(),
+        }
+        self.response_history.append(record)
+        return record
+
+
+class MCPCallTracker:
+    """
+    Tracks MCP tool calls across streaming chunks to handle deduplication.
+
+    Uses hashing to identify unique tool calls and timestamps to track when they occurred.
+    This ensures we don't double-count the same tool call appearing in multiple chunks.
+    """
+
+    def __init__(self):
+        """Initialize the tracker with empty storage."""
+        self.processed_calls = set()  # Store hashes of processed calls
+        self.call_history = []  # Store all unique calls with timestamps
+        self.last_chunk_calls = []  # Track calls from the last chunk for deduplication
+        self.dedup_window = 0.5  # Time window in seconds for deduplication
+
+    def get_call_hash(self, tool_name: str, tool_args: Dict[str, Any]) -> str:
+        """
+        Generate a unique hash for a tool call based on name and arguments.
+
+        Args:
+            tool_name: Name of the tool being called
+            tool_args: Arguments passed to the tool
+
+        Returns:
+            MD5 hash string identifying this specific call
+        """
+        # Create a deterministic string representation
+        content = f"{tool_name}:{json.dumps(tool_args, sort_keys=True)}"
+        return hashlib.md5(content.encode()).hexdigest()
+
+    def is_new_call(self, tool_name: str, tool_args: Dict[str, Any]) -> bool:
+        """
+        Check if this is a new tool call we haven't seen before.
+
+        Uses a time-window based approach: identical calls within the dedup_window
+        are considered duplicates (likely from streaming chunks), while those outside
+        the window are considered new calls (likely intentional repeated calls).
+
+        Args:
+            tool_name: Name of the tool being called
+            tool_args: Arguments passed to the tool
+
+        Returns:
+            True if this is a new call, False if we've seen it before
+        """
+        call_hash = self.get_call_hash(tool_name, tool_args)
+        current_time = time.time()
+
+        # Check if this call exists in recent history within the dedup window
+        for call in self.call_history[-10:]:  # Check last 10 calls for efficiency
+            if call.get("hash") == call_hash:
+                time_diff = current_time - call.get("timestamp", 0)
+                if time_diff < self.dedup_window:
+                    # This is likely a duplicate from streaming chunks
+                    return False
+                # If outside the window, treat as a new intentional call
+
+        # Mark as processed
+        self.processed_calls.add(call_hash)
+        return True
+
+    def add_call(self, tool_name: str, tool_args: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Add a new tool call to the history.
+
+        Args:
+            tool_name: Name of the tool being called
+            tool_args: Arguments passed to the tool
+
+        Returns:
+            Dictionary containing the call details with timestamp and hash
+        """
+        call_record = {
+            "name": tool_name,
+            "arguments": tool_args,
+            "timestamp": time.time(),
+            "hash": self.get_call_hash(tool_name, tool_args),
+            "sequence": len(self.call_history),  # Add sequence number for ordering
+        }
+        self.call_history.append(call_record)
+
+        # Clean up old history to prevent memory growth
+        if len(self.call_history) > 100:
+            self.call_history = self.call_history[-50:]
+
+        return call_record
+
+    def get_summary(self) -> str:
+        """
+        Get a summary of all tracked tool calls.
+
+        Returns:
+            Human-readable summary of tool usage
+        """
+        if not self.call_history:
+            return "No MCP tools called"
+
+        tool_names = [call["name"] for call in self.call_history]
+        unique_tools = list(dict.fromkeys(tool_names))  # Preserve order
+        return f"Used {len(self.call_history)} MCP tool calls: {', '.join(unique_tools)}"
+
+
+class MCPResponseExtractor:
+    """
+    Extracts MCP tool calls and responses from Gemini SDK stream chunks.
+
+    This class parses the internal SDK chunks to capture:
+    - function_call parts (tool invocations)
+    - function_response parts (tool results)
+    - Paired call-response data for tracking complete tool executions
+    """
+
+    def __init__(self):
+        """Initialize the extractor with empty storage."""
+        self.mcp_calls = []  # All tool calls
+        self.mcp_responses = []  # All tool responses
+        self.call_response_pairs = []  # Matched call-response pairs
+        self._pending_call = None  # Track current call awaiting response
+
+    def extract_function_call(self, function_call) -> Optional[Dict[str, Any]]:
+        """
+        Extract tool call information from SDK function_call object.
+
+        Tries multiple methods to extract data from different SDK versions:
+        1. Direct attributes (name, args)
+        2. Dictionary-like interface (get method)
+        3. __dict__ attributes
+        4. Protobuf _pb attributes
+        """
+        tool_name = None
+        tool_args = None
+
+        # Method 1: Direct attributes
+        tool_name = getattr(function_call, "name", None)
+        tool_args = getattr(function_call, "args", None)
+
+        # Method 2: Dictionary-like object
+        if tool_name is None:
+            try:
+                if hasattr(function_call, "get"):
+                    tool_name = function_call.get("name", None)
+                    tool_args = function_call.get("args", None)
+            except Exception:
+                pass
+
+        # Method 3: __dict__ inspection
+        if tool_name is None:
+            try:
+                if hasattr(function_call, "__dict__"):
+                    fc_dict = function_call.__dict__
+                    tool_name = fc_dict.get("name", None)
+                    tool_args = fc_dict.get("args", None)
+            except Exception:
+                pass
+
+        # Method 4: Protobuf _pb attribute
+        if tool_name is None:
+            try:
+                if hasattr(function_call, "_pb"):
+                    pb = function_call._pb
+                    if hasattr(pb, "name"):
+                        tool_name = pb.name
+                    if hasattr(pb, "args"):
+                        tool_args = pb.args
+            except Exception:
+                pass
+
+        if tool_name:
+            call_data = {
+                "name": tool_name,
+                "arguments": tool_args or {},
+                "timestamp": time.time(),
+                "raw": str(function_call)[:200],  # Truncate for logging
+            }
+            self.mcp_calls.append(call_data)
+            self._pending_call = call_data
+            return call_data
+
+        return None
+
+    def extract_function_response(self, function_response) -> Optional[Dict[str, Any]]:
+        """
+        Extract tool response information from SDK function_response object.
+
+        Uses same extraction methods as function_call for consistency.
+        """
+        tool_name = None
+        tool_response = None
+
+        # Method 1: Direct attributes
+        tool_name = getattr(function_response, "name", None)
+        tool_response = getattr(function_response, "response", None)
+
+        # Method 2: Dictionary-like object
+        if tool_name is None:
+            try:
+                if hasattr(function_response, "get"):
+                    tool_name = function_response.get("name", None)
+                    tool_response = function_response.get("response", None)
+            except Exception:
+                pass
+
+        # Method 3: __dict__ inspection
+        if tool_name is None:
+            try:
+                if hasattr(function_response, "__dict__"):
+                    fr_dict = function_response.__dict__
+                    tool_name = fr_dict.get("name", None)
+                    tool_response = fr_dict.get("response", None)
+            except Exception:
+                pass
+
+        # Method 4: Protobuf _pb attribute
+        if tool_name is None:
+            try:
+                if hasattr(function_response, "_pb"):
+                    pb = function_response._pb
+                    if hasattr(pb, "name"):
+                        tool_name = pb.name
+                    if hasattr(pb, "response"):
+                        tool_response = pb.response
+            except Exception:
+                pass
+
+        if tool_name:
+            response_data = {
+                "name": tool_name,
+                "response": tool_response or {},
+                "timestamp": time.time(),
+                "raw": str(function_response)[:500],  # Truncate for logging
+            }
+            self.mcp_responses.append(response_data)
+
+            # Pair with pending call if names match
+            if self._pending_call and self._pending_call["name"] == tool_name:
+                self.call_response_pairs.append(
+                    {
+                        "call": self._pending_call,
+                        "response": response_data,
+                        "duration": response_data["timestamp"] - self._pending_call["timestamp"],
+                        "paired_at": time.time(),
+                    },
+                )
+                self._pending_call = None
+
+            return response_data
+
+        return None
+
+    def get_summary(self) -> Dict[str, Any]:
+        """
+        Get a summary of all extracted MCP tool interactions.
+        """
+        return {
+            "total_calls": len(self.mcp_calls),
+            "total_responses": len(self.mcp_responses),
+            "paired_interactions": len(self.call_response_pairs),
+            "pending_call": self._pending_call is not None,
+            "tool_names": list(set(call["name"] for call in self.mcp_calls)),
+            "average_duration": (sum(pair["duration"] for pair in self.call_response_pairs) / len(self.call_response_pairs) if self.call_response_pairs else 0),
+        }
+
+    def clear(self):
+        """Clear all stored data."""
+        self.mcp_calls.clear()
+        self.mcp_responses.clear()
+        self.call_response_pairs.clear()
+        self._pending_call = None

massgen/backend/gemini_utils.py

@@ -0,0 +1,43 @@
+# -*- coding: utf-8 -*-
+"""
+Gemini-specific structured output models for coordination actions (voting and answer submission).
+"""
+
+import enum
+from typing import Optional
+
+try:
+    from pydantic import BaseModel, Field
+except ImportError:
+    BaseModel = None
+    Field = None
+
+
+class ActionType(enum.Enum):
+    """Action types for structured output."""
+
+    VOTE = "vote"
+    NEW_ANSWER = "new_answer"
+
+
+class VoteAction(BaseModel):
+    """Structured output for voting action."""
+
+    action: ActionType = Field(default=ActionType.VOTE, description="Action type")
+    agent_id: str = Field(description="Anonymous agent ID to vote for (e.g., 'agent1', 'agent2')")
+    reason: str = Field(description="Brief reason why this agent has the best answer")
+
+
+class NewAnswerAction(BaseModel):
+    """Structured output for new answer action."""
+
+    action: ActionType = Field(default=ActionType.NEW_ANSWER, description="Action type")
+    content: str = Field(description="Your improved answer. If any builtin tools like search or code execution were used, include how they are used here.")
+
+
+class CoordinationResponse(BaseModel):
+    """Structured response for coordination actions."""
+
+    action_type: ActionType = Field(description="Type of action to take")
+    vote_data: Optional[VoteAction] = Field(default=None, description="Vote data if action is vote")
+    answer_data: Optional[NewAnswerAction] = Field(default=None, description="Answer data if action is new_answer")

massgen/backend/response.py

@@ -23,16 +23,18 @@ 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_mcp import MCPBackend, UploadFileError
+from .base_with_custom_tool_and_mcp import CustomToolAndMCPBackend, UploadFileError
 
 
-class ResponseBackend(MCPBackend):
+class ResponseBackend(CustomToolAndMCPBackend):
     """Backend using the standard Response API format with multimodal support."""
 
     def __init__(self, api_key: Optional[str] = None, **kwargs):
         super().__init__(api_key, **kwargs)
         self.api_key = api_key or os.getenv("OPENAI_API_KEY")
         self.formatter = ResponseFormatter()
+
+        # Initialize API params handler after custom_tool_manager
         self.api_params_handler = ResponseAPIParamsHandler(self)
 
         # Queue for pending image saves
@@ -88,7 +90,7 @@ class ResponseBackend(MCPBackend):
                 except Exception:
                     pass
 
-    async def _stream_without_mcp_tools(
+    async def _stream_without_custom_and_mcp_tools(
         self,
         messages: List[Dict[str, Any]],
         tools: List[Dict[str, Any]],
@@ -123,6 +125,8 @@ class ResponseBackend(MCPBackend):
                     name = tool.get("function", {}).get("name") if "function" in tool else tool.get("name")
                     if name and name in self._mcp_function_names:
                         continue
+                    if name and name in self._custom_tool_names:
+                        continue
                 elif tool.get("type") == "mcp":
                     continue
                 non_mcp_tools.append(tool)
@@ -133,7 +137,7 @@ class ResponseBackend(MCPBackend):
         async for chunk in self._process_stream(stream, all_params, agent_id):
             yield chunk
 
-    async def _stream_with_mcp_tools(
+    async def _stream_with_custom_and_mcp_tools(
         self,
         current_messages: List[Dict[str, Any]],
         tools: List[Dict[str, Any]],
@@ -218,16 +222,127 @@ class ResponseBackend(MCPBackend):
 
         # Execute any captured function calls
         if captured_function_calls and response_completed:
-            # Check if any of the function calls are NOT MCP functions
-            non_mcp_functions = [call for call in captured_function_calls if call["name"] not in self._mcp_functions]
+            # 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)
 
-            if non_mcp_functions:
-                logger.info(f"Non-MCP function calls detected: {[call['name'] for call in non_mcp_functions]}. Ending MCP processing.")
+            # If there are provider calls (non-MCP, non-custom), let API handle them
+            if provider_calls:
+                logger.info(f"Provider function calls detected: {[call['name'] for call in provider_calls]}. Ending local processing.")
                 yield TextStreamChunk(type=ChunkType.DONE, source="response_api")
                 return
 
+            # Initialize for execution
+            functions_executed = False
+            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']}")
+
+                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 TextStreamChunk(
+                        type=ChunkType.CUSTOM_TOOL_STATUS,
+                        status="function_call",
+                        content=f"Arguments for Calling {call['name']}: {call['arguments']}",
+                        source=f"custom_{call['name']}",
+                    )
+
+                    yield TextStreamChunk(
+                        type=ChunkType.CUSTOM_TOOL_STATUS,
+                        status="custom_tool_error",
+                        content=f"❌ [Custom Tool Error] {error_msg}",
+                        source=f"custom_{call['name']}",
+                    )
+
+                    # 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
+
             # Check circuit breaker status before executing MCP functions
-            if not await super()._check_circuit_breaker_before_execution():
+            if mcp_calls and not await super()._check_circuit_breaker_before_execution():
                 logger.warning("All MCP servers blocked by circuit breaker")
                 yield TextStreamChunk(
                     type=ChunkType.MCP_STATUS,
@@ -238,9 +353,8 @@ class ResponseBackend(MCPBackend):
                 yield TextStreamChunk(type=ChunkType.DONE, source="response_api")
                 return
 
-            # Execute only MCP function calls
+            # Execute MCP function calls
             mcp_functions_executed = False
-            updated_messages = current_messages.copy()
 
             # Check if planning mode is enabled - block MCP tool execution during planning
             if self.is_planning_mode_enabled():
@@ -256,8 +370,6 @@ class ResponseBackend(MCPBackend):
                 return
 
             # Ensure every captured function call gets a result to prevent hanging
-            processed_call_ids = set()
-
             for call in captured_function_calls:
                 function_name = call["name"]
                 if function_name in self._mcp_functions:
@@ -367,6 +479,7 @@ class ResponseBackend(MCPBackend):
                     )
 
                     mcp_functions_executed = True
+                    functions_executed = True
 
             # Ensure all captured function calls have results to prevent hanging
             for call in captured_function_calls:
@@ -391,15 +504,14 @@ class ResponseBackend(MCPBackend):
                     mcp_functions_executed = True
 
             # Trim history after function executions to bound memory usage
-            if mcp_functions_executed:
+            if functions_executed or mcp_functions_executed:
                 updated_messages = super()._trim_message_history(updated_messages)
 
                 # Recursive call with updated messages
-                async for chunk in self._stream_with_mcp_tools(updated_messages, tools, client, **kwargs):
+                async for chunk in self._stream_with_custom_and_mcp_tools(updated_messages, tools, client, **kwargs):
                     yield chunk
             else:
-                # No MCP functions were executed, we're done
-
+                # No functions were executed, we're done
                 yield TextStreamChunk(type=ChunkType.DONE, source="response_api")
                 return
 
@@ -628,59 +740,6 @@ class ResponseBackend(MCPBackend):
         self._vector_store_ids.clear()
         self._uploaded_file_ids.clear()
 
-    # def _save_image_sync(
-    #     self,
-    #     image_data: str,
-    #     prompt: str = None,
-    #     image_format: str = "png",
-    # ) -> Optional[str]:
-    #     """
-    #     Save generated image directly to filesystem (synchronous version).
-
-    #     Args:
-    #         image_data: Base64 encoded image data
-    #         prompt: Generation prompt (used for naming)
-    #         image_format: Image format (default png)
-
-    #     Returns:
-    #         Saved file path, or None if failed
-    #     """
-    #     try:
-    #         # Use agent's filesystem workspace if available, otherwise use current working directory
-    #         if self.filesystem_manager:
-    #             workspace_path = self.filesystem_manager.get_current_workspace()
-    #         else:
-    #             workspace_path = Path.cwd()
-
-    #         # Create generated_images subdirectory path
-    #         images_dir = workspace_path
-
-    #         # Create directory if it doesn't exist
-    #         images_dir.mkdir(parents=True, exist_ok=True)
-
-    #         # Generate filename
-    #         timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-    #         if prompt:
-    #             # Clean prompt for filename
-    #             clean_prompt = "".join(c for c in prompt[:30] if c.isalnum() or c in (" ", "-", "_")).strip()
-    #             clean_prompt = clean_prompt.replace(" ", "_")
-    #             filename = f"{timestamp}_{clean_prompt}.{image_format}"
-    #         else:
-    #             filename = f"{timestamp}_generated.{image_format}"
-
-    #         file_path = images_dir / filename
-
-    #         # Decode base64 and write to file
-    #         image_bytes = base64.b64decode(image_data)
-    #         file_path.write_bytes(image_bytes)
-
-    #         logger.info(f"Image saved to: {file_path}")
-    #         return str(file_path)
-
-    #     except Exception as e:
-    #         logger.error(f"Error saving image: {e}")
-    #         return None
-
     def _convert_mcp_tools_to_openai_format(self) -> List[Dict[str, Any]]:
         """Convert MCP tools (stdio + streamable-http) to OpenAI function declarations."""
         if not self._mcp_functions: