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

massgen/tests/test_orchestration_restart.py

@@ -0,0 +1,204 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Unit tests for orchestration restart feature.
+"""
+
+import os
+import sys
+
+import pytest
+
+# Ensure project root is on sys.path
+PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+if PROJECT_ROOT not in sys.path:
+    sys.path.insert(0, PROJECT_ROOT)
+
+
+def test_coordination_config_restart_params():
+    """Test that CoordinationConfig has restart parameters."""
+    from massgen.agent_config import CoordinationConfig
+
+    config = CoordinationConfig()
+    assert hasattr(config, "max_orchestration_restarts")
+    assert config.max_orchestration_restarts == 0  # Default
+
+
+def test_agent_config_debug_final_answer():
+    """Test that AgentConfig has debug_final_answer parameter."""
+    from massgen.agent_config import AgentConfig
+
+    config = AgentConfig()
+    assert hasattr(config, "debug_final_answer")
+    assert config.debug_final_answer is None  # Default
+
+
+def test_post_evaluation_toolkit_import():
+    """Test that PostEvaluationToolkit can be imported."""
+    from massgen.tool.workflow_toolkits import PostEvaluationToolkit
+
+    assert PostEvaluationToolkit is not None
+
+
+def test_post_evaluation_tools_function():
+    """Test that get_post_evaluation_tools function exists."""
+    from massgen.tool import get_post_evaluation_tools
+
+    tools = get_post_evaluation_tools()
+    assert len(tools) == 2
+    assert tools[0]["function"]["name"] == "submit"
+    assert tools[1]["function"]["name"] == "restart_orchestration"
+
+
+def test_submit_tool_schema():
+    """Test submit tool has correct schema."""
+    from massgen.tool import get_post_evaluation_tools
+
+    tools = get_post_evaluation_tools()
+    submit_tool = tools[0]
+
+    assert submit_tool["function"]["name"] == "submit"
+    assert "confirmed" in submit_tool["function"]["parameters"]["properties"]
+    assert submit_tool["function"]["parameters"]["properties"]["confirmed"]["enum"] == [True]
+
+
+def test_restart_orchestration_tool_schema():
+    """Test restart_orchestration tool has correct schema."""
+    from massgen.tool import get_post_evaluation_tools
+
+    tools = get_post_evaluation_tools()
+    restart_tool = tools[1]
+
+    assert restart_tool["function"]["name"] == "restart_orchestration"
+    params = restart_tool["function"]["parameters"]["properties"]
+    assert "reason" in params
+    assert "instructions" in params
+    assert set(restart_tool["function"]["parameters"]["required"]) == {"reason", "instructions"}
+
+
+def test_message_templates_post_evaluation():
+    """Test that MessageTemplates has post-evaluation methods."""
+    from massgen.message_templates import MessageTemplates
+
+    templates = MessageTemplates()
+    assert hasattr(templates, "post_evaluation_system_message")
+    assert hasattr(templates, "format_restart_context")
+
+    # Test method returns strings
+    post_eval_msg = templates.post_evaluation_system_message()
+    assert isinstance(post_eval_msg, str)
+    assert "Post-Presentation Evaluation" in post_eval_msg
+
+    restart_context = templates.format_restart_context("test reason", "test instructions")
+    assert isinstance(restart_context, str)
+    assert "PREVIOUS ATTEMPT FEEDBACK" in restart_context
+
+
+def test_orchestrator_restart_state():
+    """Test that Orchestrator has restart state tracking."""
+    from massgen.agent_config import AgentConfig, CoordinationConfig
+    from massgen.orchestrator import Orchestrator
+
+    config = AgentConfig()
+    config.coordination_config = CoordinationConfig(max_orchestration_restarts=2)
+
+    orchestrator = Orchestrator(agents={}, config=config)
+
+    assert hasattr(orchestrator, "current_attempt")
+    assert hasattr(orchestrator, "max_attempts")
+    assert hasattr(orchestrator, "restart_pending")
+    assert hasattr(orchestrator, "restart_reason")
+    assert hasattr(orchestrator, "restart_instructions")
+
+    assert orchestrator.current_attempt == 0
+    assert orchestrator.max_attempts == 3  # 1 + 2 restarts
+    assert orchestrator.restart_pending is False
+
+
+def test_orchestrator_post_evaluate_method():
+    """Test that Orchestrator has post_evaluate_answer method."""
+    import inspect
+
+    from massgen.orchestrator import Orchestrator
+
+    assert hasattr(Orchestrator, "post_evaluate_answer")
+    sig = inspect.signature(Orchestrator.post_evaluate_answer)
+    assert "selected_agent_id" in sig.parameters
+    assert "final_answer" in sig.parameters
+
+
+def test_orchestrator_handle_restart_method():
+    """Test that Orchestrator has handle_restart method."""
+    from massgen.orchestrator import Orchestrator
+
+    assert hasattr(Orchestrator, "handle_restart")
+
+
+def test_handle_restart_resets_state():
+    """Test that handle_restart resets orchestrator state."""
+    from massgen.agent_config import AgentConfig, CoordinationConfig
+    from massgen.orchestrator import Orchestrator
+
+    config = AgentConfig()
+    config.coordination_config = CoordinationConfig(max_orchestration_restarts=2)
+
+    # Initialize with empty agents dict (simpler for testing)
+    orchestrator = Orchestrator(agents={}, config=config)
+
+    # Simulate state after first attempt
+    orchestrator.current_attempt = 0
+    orchestrator.restart_reason = "test reason"
+    orchestrator.restart_instructions = "test instructions"
+    orchestrator.workflow_phase = "presenting"
+    orchestrator._selected_agent = "agent1"
+    orchestrator._final_presentation_content = "some content"
+
+    # Call handle_restart
+    orchestrator.handle_restart()
+
+    # Verify state reset
+    assert orchestrator.current_attempt == 1
+    assert orchestrator.workflow_phase == "idle"
+    assert orchestrator._selected_agent is None
+    assert orchestrator._final_presentation_content is None
+    # Restart reason/instructions should be preserved for next attempt
+    assert orchestrator.restart_reason == "test reason"
+    assert orchestrator.restart_instructions == "test instructions"
+
+
+def test_base_display_restart_methods():
+    """Test that BaseDisplay has restart abstract methods."""
+    import inspect
+
+    from massgen.frontend.displays.base_display import BaseDisplay
+
+    abstract_methods = {name for name, method in inspect.getmembers(BaseDisplay, predicate=inspect.isfunction) if getattr(method, "__isabstractmethod__", False)}
+
+    assert "show_post_evaluation_content" in abstract_methods
+    assert "show_restart_banner" in abstract_methods
+    assert "show_restart_context_panel" in abstract_methods
+
+
+@pytest.mark.asyncio
+async def test_post_evaluation_tools_api_formats():
+    """Test post-evaluation tools work with different API formats."""
+    from massgen.tool import get_post_evaluation_tools
+
+    # Test chat_completions format (default)
+    tools_chat = get_post_evaluation_tools(api_format="chat_completions")
+    assert len(tools_chat) == 2
+    assert tools_chat[0]["type"] == "function"
+
+    # Test claude format
+    tools_claude = get_post_evaluation_tools(api_format="claude")
+    assert len(tools_claude) == 2
+    assert "input_schema" in tools_claude[0]
+
+    # Test response format
+    tools_response = get_post_evaluation_tools(api_format="response")
+    assert len(tools_response) == 2
+    assert tools_response[0]["type"] == "function"
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

massgen/tool/__init__.py

@@ -8,8 +8,10 @@ from ._result import ExecutionResult
 from .workflow_toolkits import (
     BaseToolkit,
     NewAnswerToolkit,
+    PostEvaluationToolkit,
     ToolType,
     VoteToolkit,
+    get_post_evaluation_tools,
     get_workflow_tools,
 )
 
@@ -35,5 +37,7 @@ __all__ = [
     "ToolType",
     "NewAnswerToolkit",
     "VoteToolkit",
+    "PostEvaluationToolkit",
     "get_workflow_tools",
+    "get_post_evaluation_tools",
 ]

massgen/tool/_multimodal_tools/understand_audio.py

@@ -0,0 +1,193 @@
+# -*- coding: utf-8 -*-
+"""
+Transcribe audio file(s) to text using OpenAI's Transcription API.
+"""
+
+import json
+import os
+from pathlib import Path
+from typing import List, Optional
+
+from dotenv import load_dotenv
+from openai import OpenAI
+
+from massgen.tool._result import ExecutionResult, TextContent
+
+
+def _validate_path_access(path: Path, allowed_paths: Optional[List[Path]] = None) -> None:
+    """
+    Validate that a path is within allowed directories.
+
+    Args:
+        path: Path to validate
+        allowed_paths: List of allowed base paths (optional)
+
+    Raises:
+        ValueError: If path is not within allowed directories
+    """
+    if not allowed_paths:
+        return  # No restrictions
+
+    for allowed_path in allowed_paths:
+        try:
+            path.relative_to(allowed_path)
+            return  # Path is within this allowed directory
+        except ValueError:
+            continue
+
+    raise ValueError(f"Path not in allowed directories: {path}")
+
+
+async def understand_audio(
+    audio_paths: List[str],
+    model: str = "gpt-4o-transcribe",
+    allowed_paths: Optional[List[str]] = None,
+) -> ExecutionResult:
+    """
+    Transcribe audio file(s) to text using OpenAI's Transcription API.
+
+    This tool processes one or more audio files through OpenAI's Transcription API
+    to extract the text content from the audio. Each file is processed separately.
+
+    Args:
+        audio_paths: List of paths to input audio files (WAV, MP3, M4A, etc.)
+                    - Relative path: Resolved relative to workspace
+                    - Absolute path: Must be within allowed directories
+        model: Model to use (default: "gpt-4o-transcribe")
+        allowed_paths: List of allowed base paths for validation (optional)
+
+    Returns:
+        ExecutionResult containing:
+        - success: Whether operation succeeded
+        - operation: "generate_text_with_input_audio"
+        - transcriptions: List of transcription results for each file
+        - audio_files: List of paths to the input audio files
+        - model: Model used
+
+    Examples:
+        generate_text_with_input_audio(["recording.wav"])
+        → Returns transcription for recording.wav
+
+        generate_text_with_input_audio(["interview1.mp3", "interview2.mp3"])
+        → Returns separate transcriptions for each file
+
+    Security:
+        - Requires valid OpenAI API key
+        - All input audio files must exist and be readable
+    """
+    try:
+        # Convert allowed_paths from strings to Path objects
+        allowed_paths_list = [Path(p) for p in allowed_paths] if allowed_paths else None
+
+        # Load environment variables
+        script_dir = Path(__file__).parent.parent.parent.parent
+        env_path = script_dir / ".env"
+        if env_path.exists():
+            load_dotenv(env_path)
+        else:
+            load_dotenv()
+
+        openai_api_key = os.getenv("OPENAI_API_KEY")
+
+        if not openai_api_key:
+            result = {
+                "success": False,
+                "operation": "generate_text_with_input_audio",
+                "error": "OpenAI API key not found. Please set OPENAI_API_KEY in .env file or environment variable.",
+            }
+            return ExecutionResult(
+                output_blocks=[TextContent(data=json.dumps(result, indent=2))],
+            )
+
+        # Initialize OpenAI client
+        client = OpenAI(api_key=openai_api_key)
+
+        # Validate and process input audio files
+        validated_audio_paths = []
+        audio_extensions = [".wav", ".mp3", ".m4a", ".mp4", ".ogg", ".flac", ".aac", ".wma", ".opus"]
+
+        for audio_path_str in audio_paths:
+            # Resolve audio path
+            if Path(audio_path_str).is_absolute():
+                audio_path = Path(audio_path_str).resolve()
+            else:
+                audio_path = (Path.cwd() / audio_path_str).resolve()
+
+            # Validate audio path
+            _validate_path_access(audio_path, allowed_paths_list)
+
+            if not audio_path.exists():
+                result = {
+                    "success": False,
+                    "operation": "generate_text_with_input_audio",
+                    "error": f"Audio file does not exist: {audio_path}",
+                }
+                return ExecutionResult(
+                    output_blocks=[TextContent(data=json.dumps(result, indent=2))],
+                )
+
+            # Check if file is an audio file
+            if audio_path.suffix.lower() not in audio_extensions:
+                result = {
+                    "success": False,
+                    "operation": "generate_text_with_input_audio",
+                    "error": f"File does not appear to be an audio file: {audio_path}",
+                }
+                return ExecutionResult(
+                    output_blocks=[TextContent(data=json.dumps(result, indent=2))],
+                )
+
+            validated_audio_paths.append(audio_path)
+
+        # Process each audio file separately using OpenAI Transcription API
+        transcriptions = []
+
+        for audio_path in validated_audio_paths:
+            try:
+                # Open audio file
+                with open(audio_path, "rb") as audio_file:
+                    # Basic transcription without prompt
+                    transcription = client.audio.transcriptions.create(
+                        model=model,
+                        file=audio_file,
+                        response_format="text",
+                    )
+
+                # Add transcription to list
+                transcriptions.append(
+                    {
+                        "file": str(audio_path),
+                        "transcription": transcription,
+                    },
+                )
+
+            except Exception as api_error:
+                result = {
+                    "success": False,
+                    "operation": "generate_text_with_input_audio",
+                    "error": f"Transcription API error for file {audio_path}: {str(api_error)}",
+                }
+                return ExecutionResult(
+                    output_blocks=[TextContent(data=json.dumps(result, indent=2))],
+                )
+
+        result = {
+            "success": True,
+            "operation": "generate_text_with_input_audio",
+            "transcriptions": transcriptions,
+            "audio_files": [str(p) for p in validated_audio_paths],
+            "model": model,
+        }
+        return ExecutionResult(
+            output_blocks=[TextContent(data=json.dumps(result, indent=2))],
+        )
+
+    except Exception as e:
+        result = {
+            "success": False,
+            "operation": "generate_text_with_input_audio",
+            "error": f"Failed to transcribe audio: {str(e)}",
+        }
+        return ExecutionResult(
+            output_blocks=[TextContent(data=json.dumps(result, indent=2))],
+        )