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

massgen/tool/_extraframework_agents/openai_assistant_lesson_planner_tool.py

@@ -0,0 +1,247 @@
+# -*- coding: utf-8 -*-
+"""
+OpenAI Assistant Lesson Planner Tool (Multi-Agent Streaming Version)
+This tool demonstrates interoperability by wrapping OpenAI's Chat Completions API with streaming support
+and multi-agent collaboration pattern similar to AG2.
+"""
+
+import os
+from typing import Any, AsyncGenerator, Dict, List
+
+from openai import AsyncOpenAI
+
+from massgen.tool import context_params
+from massgen.tool._result import ExecutionResult, TextContent
+
+# Define role-specific system prompts (similar to AG2 agents)
+CURRICULUM_AGENT_PROMPT = """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]"""
+
+LESSON_PLANNER_AGENT_PROMPT = """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."""
+
+LESSON_REVIEWER_AGENT_PROMPT = """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."""
+
+LESSON_FORMATTER_AGENT_PROMPT = """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>Detailed lesson plan activities with timing</activities>
+<assessment>Assessment details</assessment>"""
+
+
+async def run_agent_step(
+    client: AsyncOpenAI,
+    role_prompt: str,
+    user_message: str,
+    temperature: float = 0.7,
+) -> str:
+    """
+    Run a single agent step with streaming and collect the full response.
+
+    Args:
+        client: AsyncOpenAI client
+        role_prompt: System prompt for this agent role
+        user_message: User message to process
+        temperature: Temperature for generation
+
+    Returns:
+        Complete response from the agent
+    """
+    messages = [
+        {"role": "system", "content": role_prompt},
+        {"role": "user", "content": user_message},
+    ]
+
+    stream = await client.chat.completions.create(
+        model="gpt-4o",
+        messages=messages,
+        stream=True,
+        temperature=temperature,
+    )
+
+    full_response = ""
+    async for chunk in stream:
+        if chunk.choices and len(chunk.choices) > 0:
+            delta = chunk.choices[0].delta
+            if delta.content:
+                full_response += delta.content
+
+    return full_response
+
+
+@context_params("prompt")
+async def openai_assistant_lesson_planner(
+    prompt: List[Dict[str, Any]],
+) -> AsyncGenerator[ExecutionResult, None]:
+    """
+    MassGen custom tool wrapper for OpenAI lesson planner with multi-agent collaboration.
+
+    This version uses multiple specialized agents (similar to AG2) to collaboratively create
+    a lesson plan through sequential steps:
+    1. Curriculum Agent: Identifies standards and objectives
+    2. Lesson Planner Agent: Creates the detailed lesson plan
+    3. Lesson Reviewer Agent: Reviews and provides feedback
+    4. Formatter Agent: Formats the final plan
+
+    Args:
+        prompt: processed message list from orchestrator (auto-injected via execution_context)
+
+    Yields:
+        ExecutionResult containing text chunks as they arrive, or error messages
+    """
+    # 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."),
+            ],
+        )
+        return
+
+    if not prompt:
+        yield ExecutionResult(
+            output_blocks=[
+                TextContent(data="Error: No messages provided for lesson planning."),
+            ],
+        )
+        return
+
+    try:
+        # Initialize OpenAI client
+        client = AsyncOpenAI(api_key=api_key)
+
+        # Extract the user's request
+        user_request = str(prompt)
+
+        # Yield an initial message
+        yield ExecutionResult(
+            output_blocks=[
+                TextContent(data="OpenAI Lesson Planner (Multi-Agent Collaboration):\n\n"),
+            ],
+        )
+
+        # Step 1: Curriculum Agent - Determine standards and objectives
+        yield ExecutionResult(
+            output_blocks=[
+                TextContent(data="[Curriculum Agent] Identifying standards and objectives...\n"),
+            ],
+            is_log=True,
+        )
+
+        standards_and_objectives = await run_agent_step(
+            client,
+            CURRICULUM_AGENT_PROMPT,
+            f"Please provide fourth grade standards and objectives for: {user_request}",
+        )
+
+        yield ExecutionResult(
+            output_blocks=[
+                TextContent(data=f"{standards_and_objectives}\n\n"),
+            ],
+            is_log=True,
+        )
+
+        # Step 2: Lesson Planner Agent - Create detailed lesson plan
+        yield ExecutionResult(
+            output_blocks=[
+                TextContent(data="[Lesson Planner Agent] Creating detailed lesson plan...\n"),
+            ],
+            is_log=True,
+        )
+
+        lesson_plan = await run_agent_step(
+            client,
+            LESSON_PLANNER_AGENT_PROMPT,
+            f"Based on these standards and objectives:\n{standards_and_objectives}\n\nCreate a detailed lesson plan for: {user_request}",
+        )
+
+        yield ExecutionResult(
+            output_blocks=[
+                TextContent(data=f"{lesson_plan}\n\n"),
+            ],
+            is_log=True,
+        )
+
+        # Step 3: Lesson Reviewer Agent - Review and provide feedback
+        yield ExecutionResult(
+            output_blocks=[
+                TextContent(data="[Lesson Reviewer Agent] Reviewing lesson plan...\n"),
+            ],
+            is_log=True,
+        )
+
+        review_feedback = await run_agent_step(
+            client,
+            LESSON_REVIEWER_AGENT_PROMPT,
+            f"Review this lesson plan:\n{lesson_plan}\n\nProvide feedback and suggest improvements.",
+        )
+
+        yield ExecutionResult(
+            output_blocks=[
+                TextContent(data=f"{review_feedback}\n\n"),
+            ],
+            is_log=True,
+        )
+
+        # Step 4: Formatter Agent - Format the final plan with streaming
+        yield ExecutionResult(
+            output_blocks=[
+                TextContent(data="[Formatter Agent] Formatting final lesson plan...\n\n"),
+            ],
+            is_log=True,
+        )
+
+        messages = [
+            {"role": "system", "content": LESSON_FORMATTER_AGENT_PROMPT},
+            {
+                "role": "user",
+                "content": f"Format this complete lesson plan:\n\nStandards and Objectives:\n{standards_and_objectives}\n\nLesson Plan:\n{lesson_plan}\n\nReview Feedback:\n{review_feedback}",
+            },
+        ]
+
+        stream = await client.chat.completions.create(
+            model="gpt-4o",
+            messages=messages,
+            stream=True,
+            temperature=0.7,
+        )
+
+        # Stream the final formatted output
+        async for chunk in stream:
+            if chunk.choices and len(chunk.choices) > 0:
+                delta = chunk.choices[0].delta
+                if delta.content:
+                    yield ExecutionResult(
+                        output_blocks=[
+                            TextContent(data=delta.content),
+                        ],
+                    )
+
+    except Exception as e:
+        yield ExecutionResult(
+            output_blocks=[
+                TextContent(data=f"\nError during lesson planning: {str(e)}"),
+            ],
+        )

massgen/tool/_extraframework_agents/smolagent_lesson_planner_tool.py

@@ -0,0 +1,180 @@
+# -*- coding: utf-8 -*-
+"""
+SmolAgent Lesson Planner Tool
+This tool demonstrates interoperability by wrapping HuggingFace's SmolAgent framework as a MassGen custom tool.
+"""
+
+import os
+from typing import Any, AsyncGenerator, Dict, List
+
+from smolagents import CodeAgent, LiteLLMModel, tool
+
+from massgen.tool import context_params
+from massgen.tool._result import ExecutionResult, TextContent
+
+
+def run_smolagent_lesson_planner_agent(
+    messages: List[Dict[str, Any]],
+    api_key: str,
+) -> str:
+    """
+    Core SmolAgent lesson planner agent - pure SmolAgent implementation.
+
+    This function contains the pure SmolAgent logic for creating lesson plans
+    using custom tools and CodeAgent.
+
+    Args:
+        messages: Complete message history from orchestrator
+        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.")
+
+    # Define custom tools for the lesson planning workflow
+    @tool
+    def get_curriculum_standards(topic: str) -> str:
+        """
+        Determine fourth grade curriculum standards and learning objectives for a given topic.
+
+        Args:
+            topic: The lesson topic to get standards for
+
+        Returns:
+            A formatted string with standards and objectives
+        """
+        # This tool would interact with the LLM to generate standards
+        return f"Generate fourth grade curriculum standards and learning objectives for: {topic}"
+
+    @tool
+    def create_lesson_plan(topic: str, standards: str) -> str:
+        """
+        Create a detailed lesson plan based on topic and standards.
+
+        Args:
+            topic: The lesson topic
+            standards: The curriculum standards and objectives
+
+        Returns:
+            A detailed lesson plan with activities and timing
+        """
+        return f"Create a detailed lesson plan for '{topic}' based on these standards: {standards}"
+
+    @tool
+    def review_lesson_plan(lesson_plan: str) -> str:
+        """
+        Review a lesson plan for age-appropriateness, timing, and engagement.
+
+        Args:
+            lesson_plan: The lesson plan to review
+
+        Returns:
+            An improved version of the lesson plan
+        """
+        return f"Review and improve this lesson plan: {lesson_plan}"
+
+    @tool
+    def format_lesson_plan(lesson_plan: str) -> str:
+        """
+        Format a lesson plan to a standardized structure.
+
+        Args:
+            lesson_plan: The lesson plan to format
+
+        Returns:
+            A formatted lesson plan with XML-like tags
+        """
+        return (
+            f"Format this lesson plan with the following structure:\n"
+            f"<title>Lesson plan title</title>\n"
+            f"<standards>Standards covered</standards>\n"
+            f"<learning_objectives>Key learning objectives</learning_objectives>\n"
+            f"<materials>Materials required</materials>\n"
+            f"<activities>Detailed lesson plan activities with timing</activities>\n"
+            f"<assessment>Assessment details</assessment>\n\n"
+            f"Lesson plan to format: {lesson_plan}"
+        )
+
+    # Initialize the model
+    model = LiteLLMModel(
+        model_id="openai/gpt-4o",
+        api_key=api_key,
+    )
+
+    # Create the agent with custom tools
+    agent = CodeAgent(
+        tools=[get_curriculum_standards, create_lesson_plan, review_lesson_plan, format_lesson_plan],
+        model=model,
+        max_steps=10,
+    )
+
+    # Build the task from messages
+    task = f"Create a comprehensive fourth grade lesson plan for: {messages}\n\n"
+    task += "Please follow these steps:\n"
+    task += "1. Use get_curriculum_standards to identify relevant standards\n"
+    task += "2. Use create_lesson_plan to create a detailed plan\n"
+    task += "3. Use review_lesson_plan to review and improve the plan\n"
+    task += "4. Use format_lesson_plan to format the final output\n\n"
+    task += "The final plan should include:\n"
+    task += "- Opening/Hook (5-10 minutes)\n"
+    task += "- Main Activity (20-30 minutes)\n"
+    task += "- Practice Activity (15-20 minutes)\n"
+    task += "- Assessment/Closure (5-10 minutes)"
+
+    # Run the agent
+    result = agent.run(task)
+
+    return result
+
+
+@context_params("prompt")
+async def smolagent_lesson_planner(
+    prompt: List[Dict[str, Any]],
+) -> AsyncGenerator[ExecutionResult, None]:
+    """
+    MassGen custom tool wrapper for SmolAgent 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."),
+            ],
+        )
+        return
+
+    try:
+        # Call the core agent function with processed messages
+        lesson_plan = run_smolagent_lesson_planner_agent(
+            messages=prompt,
+            api_key=api_key,
+        )
+
+        yield ExecutionResult(
+            output_blocks=[
+                TextContent(data=f"SmolAgent Lesson Planner Result:\n\n{lesson_plan}"),
+            ],
+        )
+
+    except Exception as e:
+        yield ExecutionResult(
+            output_blocks=[
+                TextContent(data=f"Error creating lesson plan: {str(e)}"),
+            ],
+        )

massgen/tool/_manager.py

@@ -220,17 +220,11 @@ class ToolManager:
         if description:
             tool_schema["function"]["description"] = description
 
-        # Remove preset args from schema
-        for arg in preset_args or {}:
-            if arg in tool_schema["function"]["parameters"]["properties"]:
-                tool_schema["function"]["parameters"]["properties"].pop(arg)
-
-            if "required" in tool_schema["function"]["parameters"]:
-                if arg in tool_schema["function"]["parameters"]["required"]:
-                    tool_schema["function"]["parameters"]["required"].remove(arg)
+        # Extract context param names from decorator
+        context_param_names = getattr(base_func, "__context_params__", set())
 
-                if not tool_schema["function"]["parameters"]["required"]:
-                    tool_schema["function"]["parameters"].pop("required", None)
+        # Remove preset args and context params from schema
+        self._remove_params_from_schema(tool_schema, set(preset_args or {}) | context_param_names)
 
         tool_entry = RegisteredToolEntry(
             tool_name=tool_name,
@@ -239,6 +233,7 @@ class ToolManager:
             base_function=base_func,
             schema_def=tool_schema,
             preset_params=preset_args or {},
+            context_param_names=context_param_names,
             extension_model=None,
             post_processor=post_processor,
         )
@@ -290,11 +285,13 @@ class ToolManager:
     async def execute_tool(
         self,
         tool_request: dict,
+        execution_context: Optional[Dict[str, Any]] = None,
     ) -> AsyncGenerator[ExecutionResult, None]:
         """Execute a tool and return results as async generator.
 
         Args:
             tool_request: Tool execution request with name and input
+            execution_context: Optional execution context (messages, agent_id, etc.)
 
         Yields:
             ExecutionResult objects (accumulated)
@@ -313,13 +310,28 @@ class ToolManager:
 
         tool_entry = self.registered_tools[tool_name]
 
-        # Merge parameters: model input first, then preset params override
-        # This ensures preset_params (like agent_cwd) always take precedence
-        # and won't be overridden by null values from model
-        model_input = tool_request.get("input", {}) or {}
+        # Extract context values for marked params only
+        context_values = {}
+        if execution_context and tool_entry.context_param_names:
+            context_values = {k: v for k, v in execution_context.items() if k in tool_entry.context_param_names}
+
+        # Validate all parameters match function signature
+        self._validate_params_match_signature(
+            tool_entry.base_function,
+            tool_entry.preset_params,
+            tool_entry.context_param_names,
+            tool_request,
+            tool_name,
+        )
+
+        # Merge all parameters (validation ensures all are valid):
+        # 1. Static preset params (from registration)
+        # 2. Dynamic context values (from execution_context, marked by decorator)
+        # 3. LLM input (from tool request)
         exec_kwargs = {
-            **model_input,
-            **tool_entry.preset_params,  # preset_params override model input
+            **tool_entry.preset_params,
+            **context_values,
+            **(tool_request.get("input", {}) or {}),
         }
 
         # Prepare post-processor if exists
@@ -372,6 +384,73 @@ class ToolManager:
                 f"Tool must return ExecutionResult or Generator, got {type(result)}",
             )
 
+    @staticmethod
+    def _validate_params_match_signature(
+        func: Callable,
+        preset_params: dict,
+        context_param_names: set,
+        tool_request: dict,
+        tool_name: str,
+    ) -> None:
+        """Validate that all provided parameters match function signature.
+
+        Args:
+            func: The function to validate against
+            preset_params: Static preset parameters
+            context_param_names: Context parameter names from decorator
+            tool_request: Tool request with LLM input
+            tool_name: Tool name for error messages
+
+        Raises:
+            ValueError: If any provided parameter doesn't match function signature
+        """
+        sig = inspect.signature(func)
+        valid_params = set(sig.parameters.keys())
+
+        # Check preset args
+        invalid_preset = set(preset_params.keys()) - valid_params
+        if invalid_preset:
+            raise ValueError(
+                f"Tool '{tool_name}': preset_args contains invalid parameters: {invalid_preset}. " f"Valid parameters: {valid_params}",
+            )
+
+        # Check context params
+        invalid_context = context_param_names - valid_params
+        if invalid_context:
+            raise ValueError(
+                f"Tool '{tool_name}': @context_params decorator specifies invalid parameters: {invalid_context}. " f"Valid parameters: {valid_params}",
+            )
+
+        # Check LLM input
+        llm_input = tool_request.get("input", {}) or {}
+        invalid_llm = set(llm_input.keys()) - valid_params
+        if invalid_llm:
+            raise ValueError(
+                f"Tool '{tool_name}': LLM provided invalid parameters: {invalid_llm}. " f"Valid parameters: {valid_params}",
+            )
+
+    @staticmethod
+    def _remove_params_from_schema(tool_schema: dict, param_names: set) -> None:
+        """Remove parameters from tool schema (for preset args and context params).
+
+        Args:
+            tool_schema: The tool schema to modify
+            param_names: Set of parameter names to remove
+        """
+        for arg in param_names:
+            # Remove from properties
+            if arg in tool_schema["function"]["parameters"]["properties"]:
+                tool_schema["function"]["parameters"]["properties"].pop(arg)
+
+            # Remove from required list
+            if "required" in tool_schema["function"]["parameters"]:
+                if arg in tool_schema["function"]["parameters"]["required"]:
+                    tool_schema["function"]["parameters"]["required"].remove(arg)
+
+                # Clean up empty required list
+                if not tool_schema["function"]["parameters"]["required"]:
+                    tool_schema["function"]["parameters"].pop("required", None)
+
     def fetch_category_hints(self) -> str:
         """Get usage hints from active categories.
 
@@ -570,12 +649,19 @@ class ToolManager:
 
         func_desc = "\n\n".join(desc_parts)
 
+        # Get context param names to exclude from schema
+        context_param_names = getattr(func, "__context_params__", set())
+
         # Build parameter fields
         param_fields = {}
         for param_name, param_info in inspect.signature(func).parameters.items():
             if param_name in ["self", "cls"]:
                 continue
 
+            # Skip context params (they'll be injected at runtime)
+            if param_name in context_param_names:
+                continue
+
             if param_info.kind == inspect.Parameter.VAR_KEYWORD:
                 if not include_varkwargs:
                     continue

massgen/tool/_registered_tool.py

@@ -32,6 +32,9 @@ class RegisteredToolEntry:
     preset_params: dict[str, Any] = field(default_factory=dict)
     """Pre-configured parameters hidden from schema."""
 
+    context_param_names: set[str] = field(default_factory=set)
+    """Parameter names to inject from execution context at runtime."""
+
     extension_model: Optional[Type[BaseModel]] = None
     """Optional model for extending the base schema."""
 

massgen/tool/_result.py

@@ -59,6 +59,9 @@ class ExecutionResult:
     is_final: bool = True
     """Indicates if this is the final result in a stream."""
 
+    is_log: bool = False
+    """Indicates if this result is for logging purposes only."""
+
     was_interrupted: bool = False
     """Indicates if the execution was interrupted."""