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

massgen/tests/test_langgraph_lesson_planner.py

@@ -0,0 +1,223 @@
+# -*- coding: utf-8 -*-
+"""
+Test LangGraph Lesson Planner Tool
+Tests the interoperability feature where LangGraph state graph is wrapped as a MassGen custom tool.
+"""
+
+import asyncio
+import os
+import sys
+from pathlib import Path
+
+import pytest
+
+# Add parent directory to path for imports
+sys.path.insert(0, str(Path(__file__).parent.parent.parent))
+
+from massgen.tool._extraframework_agents.langgraph_lesson_planner_tool import (  # noqa: E402
+    langgraph_lesson_planner,
+)
+from massgen.tool._result import ExecutionResult  # noqa: E402
+
+
+class TestLangGraphLessonPlannerTool:
+    """Test LangGraph Lesson Planner Tool functionality."""
+
+    @pytest.mark.asyncio
+    async def test_basic_lesson_plan_creation(self):
+        """Test basic lesson plan creation with a simple topic."""
+        # Skip if no API key
+        api_key = os.getenv("OPENAI_API_KEY")
+        if not api_key:
+            pytest.skip("OPENAI_API_KEY not set")
+
+        # Test with a simple topic
+        result = await langgraph_lesson_planner(topic="photosynthesis", api_key=api_key)
+
+        # Verify result structure
+        assert isinstance(result, ExecutionResult)
+        assert len(result.output_blocks) > 0
+        # Check that the result doesn't contain an error
+        assert not result.output_blocks[0].data.startswith("Error:")
+
+        # Verify lesson plan contains expected elements
+        lesson_plan = result.output_blocks[0].data
+        assert "photosynthesis" in lesson_plan.lower()
+
+    @pytest.mark.asyncio
+    async def test_lesson_plan_with_env_api_key(self):
+        """Test lesson plan creation using environment variable for API key."""
+        # Skip if no API key
+        if not os.getenv("OPENAI_API_KEY"):
+            pytest.skip("OPENAI_API_KEY not set")
+
+        # Test without passing api_key parameter (should use env var)
+        result = await langgraph_lesson_planner(topic="fractions")
+
+        assert isinstance(result, ExecutionResult)
+        assert len(result.output_blocks) > 0
+        # Check that the result doesn't contain an error
+        assert not result.output_blocks[0].data.startswith("Error:")
+
+    @pytest.mark.asyncio
+    async def test_missing_api_key_error(self):
+        """Test error handling when API key is missing."""
+        # Temporarily save and remove env var
+        original_key = os.environ.get("OPENAI_API_KEY")
+        if "OPENAI_API_KEY" in os.environ:
+            del os.environ["OPENAI_API_KEY"]
+
+        try:
+            result = await langgraph_lesson_planner(topic="test topic")
+
+            # Should return error result
+            assert isinstance(result, ExecutionResult)
+            assert result.output_blocks[0].data.startswith("Error:")
+            assert "OPENAI_API_KEY not found" in result.output_blocks[0].data
+        finally:
+            # Restore env var
+            if original_key:
+                os.environ["OPENAI_API_KEY"] = original_key
+
+    @pytest.mark.asyncio
+    async def test_different_topics(self):
+        """Test lesson plan creation with different topics."""
+        # Skip if no API key
+        api_key = os.getenv("OPENAI_API_KEY")
+        if not api_key:
+            pytest.skip("OPENAI_API_KEY not set")
+
+        topics = ["addition", "animals", "water cycle"]
+
+        for topic in topics:
+            result = await langgraph_lesson_planner(topic=topic, api_key=api_key)
+
+            assert isinstance(result, ExecutionResult)
+            assert len(result.output_blocks) > 0
+            # Check that the result doesn't contain an error
+            assert not result.output_blocks[0].data.startswith("Error:")
+            assert topic.lower() in result.output_blocks[0].data.lower()
+
+    @pytest.mark.asyncio
+    async def test_concurrent_lesson_plan_creation(self):
+        """Test creating multiple lesson plans concurrently."""
+        # Skip if no API key
+        api_key = os.getenv("OPENAI_API_KEY")
+        if not api_key:
+            pytest.skip("OPENAI_API_KEY not set")
+
+        topics = ["math", "science", "reading"]
+
+        # Create tasks for concurrent execution
+        tasks = [langgraph_lesson_planner(topic=topic, api_key=api_key) for topic in topics]
+
+        # Execute concurrently
+        results = await asyncio.gather(*tasks)
+
+        # Verify all results
+        assert len(results) == len(topics)
+        for i, result in enumerate(results):
+            assert isinstance(result, ExecutionResult)
+            assert len(result.output_blocks) > 0
+            # Check that the result doesn't contain an error
+            assert not result.output_blocks[0].data.startswith("Error:")
+            assert topics[i].lower() in result.output_blocks[0].data.lower()
+
+
+class TestLangGraphToolIntegration:
+    """Test LangGraph tool integration with MassGen tool system."""
+
+    def test_tool_function_signature(self):
+        """Test that the tool has the correct async signature."""
+        import inspect
+
+        assert inspect.iscoroutinefunction(langgraph_lesson_planner)
+
+        # Get function signature
+        sig = inspect.signature(langgraph_lesson_planner)
+        params = sig.parameters
+
+        # Verify parameters
+        assert "topic" in params
+        assert "api_key" in params
+
+        # Verify return annotation
+        assert sig.return_annotation == ExecutionResult
+
+    @pytest.mark.asyncio
+    async def test_execution_result_structure(self):
+        """Test that the returned ExecutionResult has the correct structure."""
+        # Skip if no API key
+        api_key = os.getenv("OPENAI_API_KEY")
+        if not api_key:
+            pytest.skip("OPENAI_API_KEY not set")
+
+        result = await langgraph_lesson_planner(topic="test", api_key=api_key)
+
+        # Verify ExecutionResult structure
+        assert hasattr(result, "output_blocks")
+        assert isinstance(result.output_blocks, list)
+        assert len(result.output_blocks) > 0
+        # Check that the result doesn't contain an error
+        assert not result.output_blocks[0].data.startswith("Error:")
+
+        # Verify TextContent structure
+        from massgen.tool._result import TextContent
+
+        assert isinstance(result.output_blocks[0], TextContent)
+        assert hasattr(result.output_blocks[0], "data")
+        assert isinstance(result.output_blocks[0].data, str)
+
+
+class TestLangGraphToolWithBackend:
+    """Test LangGraph tool with ResponseBackend."""
+
+    @pytest.mark.asyncio
+    async def test_backend_registration(self):
+        """Test registering LangGraph tool with ResponseBackend."""
+        from massgen.backend.response import ResponseBackend
+
+        api_key = os.getenv("OPENAI_API_KEY", "test-key")
+
+        # Import the tool
+        from massgen.tool._extraframework_agents.langgraph_lesson_planner_tool import (
+            langgraph_lesson_planner,
+        )
+
+        # Register with backend
+        backend = ResponseBackend(
+            api_key=api_key,
+            custom_tools=[
+                {
+                    "func": langgraph_lesson_planner,
+                    "description": "Create a comprehensive lesson plan using LangGraph state graph",
+                },
+            ],
+        )
+
+        # Verify tool is registered
+        assert "langgraph_lesson_planner" in backend._custom_tool_names
+
+        # Verify schema generation
+        schemas = backend._get_custom_tools_schemas()
+        assert len(schemas) >= 1
+
+        # Find our tool's schema
+        langgraph_schema = None
+        for schema in schemas:
+            if schema["function"]["name"] == "langgraph_lesson_planner":
+                langgraph_schema = schema
+                break
+
+        assert langgraph_schema is not None
+        assert langgraph_schema["type"] == "function"
+        assert "parameters" in langgraph_schema["function"]
+
+
+# ============================================================================
+# Run tests
+# ============================================================================
+
+if __name__ == "__main__":
+    # Run pytest
+    pytest.main([__file__, "-v"])

massgen/tool/__init__.py

@@ -2,6 +2,7 @@
 """Tool module for MassGen framework."""
 
 from ._code_executors import run_python_script, run_shell_script
+from ._decorators import context_params
 from ._file_handlers import append_file_content, read_file_content, save_file_content
 from ._manager import ToolManager
 from ._result import ExecutionResult
@@ -18,21 +19,13 @@ from .workflow_toolkits import (
 __all__ = [
     "ToolManager",
     "ExecutionResult",
+    "context_params",
     "two_num_tool",
     "run_python_script",
     "run_shell_script",
     "read_file_content",
     "save_file_content",
     "append_file_content",
-    "dashscope_generate_image",
-    "dashscope_generate_audio",
-    "dashscope_analyze_image",
-    "openai_generate_image",
-    "openai_generate_audio",
-    "openai_modify_image",
-    "openai_create_variation",
-    "openai_analyze_image",
-    "openai_transcribe_audio",
     "BaseToolkit",
     "ToolType",
     "NewAnswerToolkit",

massgen/tool/_decorators.py

@@ -0,0 +1,52 @@
+# -*- coding: utf-8 -*-
+"""Decorators for custom tool functions."""
+
+from typing import Callable
+
+
+def context_params(*param_names: str) -> Callable[[Callable], Callable]:
+    """Mark parameters for auto-injection from ExecutionContext.
+
+    Parameters marked with this decorator will be:
+    1. Excluded from the LLM schema (not visible to the model)
+    2. Automatically injected from execution_context at runtime
+
+    This is useful for parameters that should come from the backend runtime
+    context rather than from the LLM, such as:
+    - messages: Conversation history
+    - agent_id: Current agent identifier
+    - backend_name: Backend provider name
+    - current_stage: Current coordination stage
+
+    Args:
+        *param_names: Names of parameters to mark as context parameters
+
+    Returns:
+        Decorator function that marks the parameters
+
+    Example:
+        >>> from massgen.tool import context_params, ExecutionResult
+        >>> from typing import List, Dict, Any
+        >>>
+        >>> @context_params("messages", "agent_id")
+        >>> async def analyze_conversation(
+        ...     query: str,  # LLM provides this
+        ...     messages: List[Dict[str, Any]],  # Auto-injected from context
+        ...     agent_id: str,  # Auto-injected from context
+        ... ) -> ExecutionResult:
+        ...     '''Analyze conversation with full context.'''
+        ...     # messages and agent_id are automatically filled from execution_context
+        ...     system_msg = next((m for m in messages if m.get("role") == "system"), None)
+        ...     return ExecutionResult(...)
+
+    Note:
+        The execution_context is provided by the backend when executing tools.
+        Only parameters marked by this decorator will be injected from the context.
+    """
+
+    def decorator(func: Callable) -> Callable:
+        """Store context parameter names in function metadata."""
+        func.__context_params__ = set(param_names)
+        return func
+
+    return decorator

massgen/tool/_extraframework_agents/ag2_lesson_planner_tool.py

@@ -0,0 +1,251 @@
+# -*- coding: utf-8 -*-
+"""
+AG2 (AutoGen) Nested Chat Lesson Planner Tool
+This tool demonstrates interoperability by wrapping AutoGen's nested chat functionality as a MassGen custom tool.
+"""
+
+import os
+from typing import Any, AsyncGenerator, Dict, List
+
+from autogen import ConversableAgent, GroupChat, GroupChatManager
+
+from massgen.tool import context_params
+from massgen.tool._result import ExecutionResult, TextContent
+
+
+def run_ag2_lesson_planner_agent(
+    messages: List[Dict[str, Any]],
+    api_key: str,
+):
+    """
+    Core AG2 lesson planner agent - pure AutoGen implementation.
+
+    This function contains the pure AG2/AutoGen logic for creating lesson plans
+    using nested chats and multiple specialized agents.
+
+    Args:
+        messages: messages for the agent to execute
+        api_key: OpenAI API key for the agents
+
+    Returns:
+        The formatted lesson plan as a string
+
+    Raises:
+        Exception: Any errors during agent execution
+    """
+    if not messages:
+        raise ValueError("No messages provided for lesson planning.")
+    # Configure LLM
+    llm_config = {
+        "config_list": [
+            {
+                "api_type": "openai",
+                "model": "gpt-4o",
+                "api_key": api_key,
+            },
+        ],
+    }
+
+    # Curriculum Standards Agent
+    curriculum_agent = ConversableAgent(
+        name="Curriculum_Agent",
+        system_message="""You are a curriculum standards expert for fourth grade education.
+        When given a topic, you provide relevant grade-level standards and learning objectives.
+        Format every response as:
+        STANDARDS:
+        - [Standard 1]
+        - [Standard 2]
+        OBJECTIVES:
+        - By the end of this lesson, students will be able to [objective 1]
+        - By the end of this lesson, students will be able to [objective 2]""",
+        human_input_mode="NEVER",
+        llm_config=llm_config,
+    )
+
+    # Lesson Planner Agent
+    lesson_planner_agent = ConversableAgent(
+        name="Lesson_Planner_Agent",
+        system_message="""You are a lesson planning specialist.
+        Given standards and objectives, you create detailed lesson plans including:
+        - Opening/Hook (5-10 minutes)
+        - Main Activity (20-30 minutes)
+        - Practice Activity (15-20 minutes)
+        - Assessment/Closure (5-10 minutes)
+        Format as a structured lesson plan with clear timing and materials needed.""",
+        human_input_mode="NEVER",
+        llm_config=llm_config,
+    )
+
+    # Lesson Reviewer Agent
+    lesson_reviewer_agent = ConversableAgent(
+        name="Lesson_Reviewer_Agent",
+        system_message="""You are a lesson plan reviewer who ensures:
+        1. Age-appropriate content and activities
+        2. Alignment with provided standards
+        3. Realistic timing
+        4. Clear instructions
+        5. Differentiation opportunities
+        Provide specific feedback in these areas and suggest improvements if needed.""",
+        human_input_mode="NEVER",
+        llm_config=llm_config,
+    )
+
+    # Lead Teacher Agent
+    lead_teacher_agent = ConversableAgent(
+        name="Lead_Teacher_Agent",
+        system_message="""You are an experienced fourth grade teacher who oversees the lesson planning process.
+        Your role is to:
+        1. Initiate the planning process with a clear topic
+        2. Review and integrate feedback from other agents
+        3. Ensure the final lesson plan is practical and engaging
+        4. Make final adjustments based on classroom experience""",
+        human_input_mode="NEVER",
+        llm_config=llm_config,
+    )
+
+    # Create the group chat for collaborative lesson planning
+    planning_chat = GroupChat(
+        agents=[curriculum_agent, lesson_planner_agent, lesson_reviewer_agent],
+        messages=[],
+        max_round=2,
+        send_introductions=True,
+    )
+
+    planning_manager = GroupChatManager(
+        groupchat=planning_chat,
+        llm_config=llm_config,
+    )
+
+    # Formatter of the final lesson plan to a standard format
+    formatter_message = """You are a lesson plan formatter. Format the complete plan as follows:
+<title>Lesson plan title</title>
+<standards>Standards covered</standards>
+<learning_objectives>Key learning objectives</learning_objectives>
+<materials>Materials required</materials>
+<activities>Lesson plan activities</activities>
+<assessment>Assessment details</assessment>
+"""
+
+    lesson_formatter = ConversableAgent(
+        name="formatter_agent",
+        system_message=formatter_message,
+        llm_config=llm_config,
+    )
+
+    # Create nested chats configuration
+    nested_chats = [
+        {
+            # The first internal chat determines the standards and objectives
+            "recipient": curriculum_agent,
+            "message": f"Please provide fourth grade standards and objectives for: {messages}",
+            "max_turns": 2,
+            "summary_method": "last_msg",
+        },
+        {
+            # A group chat follows, where the lesson plan is created
+            "recipient": planning_manager,
+            "message": "Based on these standards and objectives, create a detailed lesson plan.",
+            "max_turns": 1,
+            "summary_method": "last_msg",
+        },
+        {
+            # Finally, a two-agent chat formats the lesson plan
+            "recipient": lesson_formatter,
+            "message": "Format the lesson plan.",
+            "max_turns": 1,
+            "summary_method": "last_msg",
+        },
+    ]
+
+    # Register nested chats with the lead teacher
+    lead_teacher_agent.register_nested_chats(
+        chat_queue=nested_chats,
+        trigger=lambda sender: sender.name == "Assistant_Agent",
+    )
+
+    # Create a simple assistant agent to trigger the nested chat
+    assistant_agent = ConversableAgent(
+        name="Assistant_Agent",
+        system_message="You are a helpful assistant.",
+        human_input_mode="NEVER",
+        llm_config=llm_config,
+    )
+
+    # Initiate the chat and get the result
+    response = assistant_agent.run(
+        recipient=lead_teacher_agent,
+        message=str(messages),
+        max_turns=1,
+    )
+
+    # Extract the lesson plan from the result
+
+    return response
+
+
+@context_params("prompt")
+async def ag2_lesson_planner(
+    prompt: List[Dict[str, Any]],
+) -> AsyncGenerator[ExecutionResult, None]:
+    """
+    MassGen custom tool wrapper for AG2 lesson planner.
+
+    This is the interface exposed to MassGen's backend. It handles environment setup,
+    error handling, and wraps the core agent logic in ExecutionResult.
+
+    Args:
+        prompt: processed message list from orchestrator (auto-injected via execution_context)
+
+    Returns:
+        ExecutionResult containing the formatted lesson plan or error message
+    """
+    # Get API key from environment
+    api_key = os.getenv("OPENAI_API_KEY")
+
+    if not api_key:
+        yield ExecutionResult(
+            output_blocks=[
+                TextContent(data="Error: OPENAI_API_KEY not found. Please set the environment variable."),
+            ],
+        )
+
+    try:
+        # Call the core agent function with processed messages
+        response = run_ag2_lesson_planner_agent(
+            messages=prompt,
+            api_key=api_key,
+        )
+
+        last_nested_chat_event_msgs = []
+
+        def process_and_log_event(*args, **kwargs) -> None:
+            """Process and log AG2 event, returning string representation."""
+            line = " ".join(str(arg) for arg in args)
+            last_nested_chat_event_msgs.append(line)
+
+        for event in response.events:
+            last_nested_chat_event_msgs.clear()
+            event.print(f=process_and_log_event)
+            formatted_message = "\n".join(last_nested_chat_event_msgs)
+
+            if event.type == "run_completion":
+                # Final output
+                yield ExecutionResult(
+                    output_blocks=[
+                        TextContent(data=event.content.summary),
+                    ],
+                )
+            else:
+                yield ExecutionResult(
+                    output_blocks=[
+                        TextContent(data=formatted_message),
+                    ],
+                    is_log=True,
+                )
+
+    except Exception as e:
+        yield ExecutionResult(
+            output_blocks=[
+                TextContent(data=f"Error creating lesson plan: {str(e)}"),
+            ],
+        )