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

massgen/formatter/_gemini_formatter.py

@@ -75,6 +75,21 @@ class GeminiFormatter(FormatterBase):
 
         return "vote" in tool_names and "new_answer" in tool_names
 
+    def has_post_evaluation_tools(self, tools: List[Dict[str, Any]]) -> bool:
+        """Detect if tools contain submit/restart_orchestration post-evaluation tools."""
+        if not tools:
+            return False
+
+        tool_names = set()
+        for tool in tools:
+            if tool.get("type") == "function":
+                if "function" in tool:
+                    tool_names.add(tool["function"].get("name", ""))
+                elif "name" in tool:
+                    tool_names.add(tool.get("name", ""))
+
+        return "submit" in tool_names and "restart_orchestration" in tool_names
+
     def build_structured_output_prompt(self, base_content: str, valid_agent_ids: Optional[List[str]] = None) -> str:
         """Build prompt that encourages structured output for coordination."""
         agent_list = ""
@@ -104,6 +119,33 @@ If you want to provide a NEW ANSWER:
   }}
 }}
 
+Make your decision and include the JSON at the very end of your response."""
+
+    def build_post_evaluation_prompt(self, base_content: str) -> str:
+        """Build prompt that encourages structured output for post-evaluation."""
+        return f"""{base_content}
+
+IMPORTANT: You must respond with a structured JSON decision at the end of your response.
+
+If you want to SUBMIT the answer (it's complete and satisfactory):
+{{
+  "action_type": "submit",
+  "submit_data": {{
+    "action": "submit",
+    "confirmed": true
+  }}
+}}
+
+If you want to RESTART with improvements:
+{{
+  "action_type": "restart",
+  "restart_data": {{
+    "action": "restart",
+    "reason": "Clear explanation of why the answer is insufficient",
+    "instructions": "Detailed, actionable guidance for the next attempt"
+  }}
+}}
+
 Make your decision and include the JSON at the very end of your response."""
 
     def extract_structured_response(self, response_text: str) -> Optional[Dict[str, Any]]:
@@ -192,6 +234,7 @@ Make your decision and include the JSON at the very end of your response."""
         """Convert structured response to tool call format."""
         action_type = structured_response.get("action_type")
 
+        # Coordination tools
         if action_type == "vote":
             vote_data = structured_response.get("vote_data", {})
             return [
@@ -221,6 +264,36 @@ Make your decision and include the JSON at the very end of your response."""
                 },
             ]
 
+        # Post-evaluation tools
+        elif action_type == "submit":
+            submit_data = structured_response.get("submit_data", {})
+            return [
+                {
+                    "id": f"submit_{abs(hash(str(submit_data))) % 10000 + 1}",
+                    "type": "function",
+                    "function": {
+                        "name": "submit",
+                        "arguments": {"confirmed": submit_data.get("confirmed", True)},
+                    },
+                },
+            ]
+
+        elif action_type == "restart":
+            restart_data = structured_response.get("restart_data", {})
+            return [
+                {
+                    "id": f"restart_{abs(hash(str(restart_data))) % 10000 + 1}",
+                    "type": "function",
+                    "function": {
+                        "name": "restart_orchestration",
+                        "arguments": {
+                            "reason": restart_data.get("reason", ""),
+                            "instructions": restart_data.get("instructions", ""),
+                        },
+                    },
+                },
+            ]
+
         return []
 
     # Custom tools formatting for Gemini