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

massgen/tool/workflow_toolkits/new_answer.py

@@ -0,0 +1,126 @@
+# -*- coding: utf-8 -*-
+"""
+New Answer toolkit for MassGen workflow coordination.
+"""
+
+from typing import Any, Dict, List, Optional
+
+from .base import BaseToolkit, ToolType
+
+
+class NewAnswerToolkit(BaseToolkit):
+    """New Answer toolkit for agent coordination workflows."""
+
+    def __init__(self, template_overrides: Optional[Dict[str, Any]] = None):
+        """
+        Initialize the New Answer toolkit.
+
+        Args:
+            template_overrides: Optional template overrides for customization
+        """
+        self._template_overrides = template_overrides or {}
+
+    @property
+    def toolkit_id(self) -> str:
+        """Unique identifier for new answer toolkit."""
+        return "new_answer"
+
+    @property
+    def toolkit_type(self) -> ToolType:
+        """Type of this toolkit."""
+        return ToolType.WORKFLOW
+
+    def is_enabled(self, config: Dict[str, Any]) -> bool:
+        """
+        Check if new answer is enabled in configuration.
+
+        Args:
+            config: Configuration dictionary.
+
+        Returns:
+            True if workflow tools are enabled or not explicitly disabled.
+        """
+        # Enable by default for workflow, unless explicitly disabled
+        return config.get("enable_workflow_tools", True)
+
+    def get_tools(self, config: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """
+        Get new answer tool definition based on API format.
+
+        Args:
+            config: Configuration including api_format.
+
+        Returns:
+            List containing the new answer tool definition.
+        """
+        # Check for template override
+        if "new_answer_tool" in self._template_overrides:
+            override = self._template_overrides["new_answer_tool"]
+            if callable(override):
+                return [override()]
+            return [override]
+
+        api_format = config.get("api_format", "chat_completions")
+
+        if api_format == "claude":
+            # Claude native format
+            return [
+                {
+                    "name": "new_answer",
+                    "description": "Submit a new and improved answer",
+                    "input_schema": {
+                        "type": "object",
+                        "properties": {
+                            "content": {
+                                "type": "string",
+                                "description": "Your improved answer. If any builtin tools like search or code execution were used, mention how they are used here.",
+                            },
+                        },
+                        "required": ["content"],
+                    },
+                },
+            ]
+
+        elif api_format == "response":
+            # Response API format (OpenAI-style but simpler)
+            return [
+                {
+                    "type": "function",
+                    "function": {
+                        "name": "new_answer",
+                        "description": "Submit a new and improved answer",
+                        "parameters": {
+                            "type": "object",
+                            "properties": {
+                                "content": {
+                                    "type": "string",
+                                    "description": "Your improved answer. If any builtin tools like search or code execution were used, mention how they are used here.",
+                                },
+                            },
+                            "required": ["content"],
+                        },
+                    },
+                },
+            ]
+
+        else:
+            # Default Chat Completions format
+            return [
+                {
+                    "type": "function",
+                    "function": {
+                        "name": "new_answer",
+                        "description": "Submit a new and improved answer",
+                        "parameters": {
+                            "type": "object",
+                            "properties": {
+                                "content": {
+                                    "type": "string",
+                                    "description": "Your improved answer. If any builtin tools like search or code execution were used, mention how they are used here.",
+                                },
+                            },
+                            "required": ["content"],
+                        },
+                    },
+                },
+            ]

massgen/tool/workflow_toolkits/vote.py

@@ -0,0 +1,167 @@
+# -*- coding: utf-8 -*-
+"""
+Vote toolkit for MassGen workflow coordination.
+"""
+
+from typing import Any, Dict, List, Optional
+
+from .base import BaseToolkit, ToolType
+
+
+class VoteToolkit(BaseToolkit):
+    """Vote toolkit for agent coordination workflows."""
+
+    def __init__(self, valid_agent_ids: Optional[List[str]] = None, template_overrides: Optional[Dict[str, Any]] = None):
+        """
+        Initialize the Vote toolkit.
+
+        Args:
+            valid_agent_ids: List of valid agent IDs for voting
+            template_overrides: Optional template overrides for customization
+        """
+        self.valid_agent_ids = valid_agent_ids
+        self._template_overrides = template_overrides or {}
+
+    @property
+    def toolkit_id(self) -> str:
+        """Unique identifier for vote toolkit."""
+        return "vote"
+
+    @property
+    def toolkit_type(self) -> ToolType:
+        """Type of this toolkit."""
+        return ToolType.WORKFLOW
+
+    def is_enabled(self, config: Dict[str, Any]) -> bool:
+        """
+        Check if vote is enabled in configuration.
+
+        Args:
+            config: Configuration dictionary.
+
+        Returns:
+            True if workflow tools are enabled or not explicitly disabled.
+        """
+        # Enable by default for workflow, unless explicitly disabled
+        return config.get("enable_workflow_tools", True)
+
+    def set_valid_agent_ids(self, agent_ids: List[str]):
+        """
+        Update the valid agent IDs for voting.
+
+        Args:
+            agent_ids: List of valid agent IDs
+        """
+        self.valid_agent_ids = agent_ids
+
+    def get_tools(self, config: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """
+        Get vote tool definition based on API format.
+
+        Args:
+            config: Configuration including api_format and potentially valid_agent_ids.
+
+        Returns:
+            List containing the vote tool definition.
+        """
+        # Check for template override
+        if "vote_tool" in self._template_overrides:
+            override = self._template_overrides["vote_tool"]
+            if callable(override):
+                return [override(self.valid_agent_ids)]
+            return [override]
+
+        # Get valid agent IDs from config if not set
+        valid_ids = config.get("valid_agent_ids", self.valid_agent_ids)
+
+        api_format = config.get("api_format", "chat_completions")
+
+        if api_format == "claude":
+            # Claude native format
+            tool_def = {
+                "name": "vote",
+                "description": "Vote for the best agent to present final answer",
+                "input_schema": {
+                    "type": "object",
+                    "properties": {
+                        "agent_id": {
+                            "type": "string",
+                            "description": "Anonymous agent ID to vote for (e.g., 'agent1', 'agent2')",
+                        },
+                        "reason": {
+                            "type": "string",
+                            "description": "Brief reason why this agent has the best answer",
+                        },
+                    },
+                    "required": ["agent_id", "reason"],
+                },
+            }
+
+            # Add enum constraint for valid agent IDs
+            if valid_ids:
+                anon_agent_ids = [f"agent{i}" for i in range(1, len(valid_ids) + 1)]
+                tool_def["input_schema"]["properties"]["agent_id"]["enum"] = anon_agent_ids
+
+            return [tool_def]
+
+        elif api_format == "response":
+            # Response API format
+            tool_def = {
+                "type": "function",
+                "function": {
+                    "name": "vote",
+                    "description": "Vote for the best agent to present final answer",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "agent_id": {
+                                "type": "string",
+                                "description": "Anonymous agent ID to vote for (e.g., 'agent1', 'agent2')",
+                            },
+                            "reason": {
+                                "type": "string",
+                                "description": "Brief reason why this agent has the best answer",
+                            },
+                        },
+                        "required": ["agent_id", "reason"],
+                    },
+                },
+            }
+
+            # Add enum constraint for valid agent IDs
+            if valid_ids:
+                anon_agent_ids = [f"agent{i}" for i in range(1, len(valid_ids) + 1)]
+                tool_def["function"]["parameters"]["properties"]["agent_id"]["enum"] = anon_agent_ids
+
+            return [tool_def]
+
+        else:
+            # Default Chat Completions format
+            tool_def = {
+                "type": "function",
+                "function": {
+                    "name": "vote",
+                    "description": "Vote for the best agent to present final answer",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "agent_id": {
+                                "type": "string",
+                                "description": "Anonymous agent ID to vote for (e.g., 'agent1', 'agent2')",
+                            },
+                            "reason": {
+                                "type": "string",
+                                "description": "Brief reason why this agent has the best answer",
+                            },
+                        },
+                        "required": ["agent_id", "reason"],
+                    },
+                },
+            }
+
+            # Add enum constraint for valid agent IDs
+            if valid_ids:
+                anon_agent_ids = [f"agent{i}" for i in range(1, len(valid_ids) + 1)]
+                tool_def["function"]["parameters"]["properties"]["agent_id"]["enum"] = anon_agent_ids
+
+            return [tool_def]