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

massgen/tool/__init__.py

@@ -0,0 +1,39 @@
+# -*- coding: utf-8 -*-
+"""Tool module for MassGen framework."""
+
+from ._code_executors import run_python_script, run_shell_script
+from ._file_handlers import append_file_content, read_file_content, save_file_content
+from ._manager import ToolManager
+from ._result import ExecutionResult
+from .workflow_toolkits import (
+    BaseToolkit,
+    NewAnswerToolkit,
+    ToolType,
+    VoteToolkit,
+    get_workflow_tools,
+)
+
+__all__ = [
+    "ToolManager",
+    "ExecutionResult",
+    "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",
+    "VoteToolkit",
+    "get_workflow_tools",
+]

massgen/tool/_async_helpers.py

@@ -0,0 +1,70 @@
+# -*- coding: utf-8 -*-
+"""Async utility functions for wrapping different return types."""
+
+import asyncio
+from typing import AsyncGenerator, Callable, Generator, Optional
+
+from ._result import ExecutionResult, TextContent
+
+
+async def _apply_post_processing(
+    result: ExecutionResult,
+    processor: Optional[Callable[[ExecutionResult], Optional[ExecutionResult]]],
+) -> ExecutionResult:
+    """Apply post-processing to an execution result."""
+    if processor:
+        processed = processor(result)
+        if processed:
+            return processed
+    return result
+
+
+async def wrap_object_async(
+    obj: ExecutionResult,
+    processor: Optional[Callable[[ExecutionResult], Optional[ExecutionResult]]],
+) -> AsyncGenerator[ExecutionResult, None]:
+    """Convert a single ExecutionResult to async generator."""
+    yield await _apply_post_processing(obj, processor)
+
+
+async def wrap_sync_gen_async(
+    sync_gen: Generator[ExecutionResult, None, None],
+    processor: Optional[Callable[[ExecutionResult], Optional[ExecutionResult]]],
+) -> AsyncGenerator[ExecutionResult, None]:
+    """Convert sync generator to async generator."""
+    for chunk in sync_gen:
+        yield await _apply_post_processing(chunk, processor)
+
+
+async def wrap_as_async_generator(
+    async_gen: AsyncGenerator[ExecutionResult, None],
+    processor: Optional[Callable[[ExecutionResult], Optional[ExecutionResult]]],
+) -> AsyncGenerator[ExecutionResult, None]:
+    """Wrap async generator with interruption handling."""
+
+    previous_chunk = None
+    try:
+        async for chunk in async_gen:
+            processed = await _apply_post_processing(chunk, processor)
+            yield processed
+            previous_chunk = processed
+
+    except asyncio.CancelledError:
+        interrupt_msg = TextContent(
+            data="<system>Execution interrupted by user request</system>",
+        )
+
+        if previous_chunk:
+            previous_chunk.output_blocks.append(interrupt_msg)
+            previous_chunk.was_interrupted = True
+            previous_chunk.is_final = True
+            yield await _apply_post_processing(previous_chunk, processor)
+        else:
+            yield await _apply_post_processing(
+                ExecutionResult(
+                    output_blocks=[interrupt_msg],
+                    was_interrupted=True,
+                    is_final=True,
+                ),
+                processor,
+            )

massgen/tool/_basic/__init__.py

@@ -0,0 +1,8 @@
+# -*- coding: utf-8 -*-
+"""Basic tools."""
+
+from ._two_num_tool import two_num_tool
+
+__all__ = [
+    "two_num_tool",
+]

massgen/tool/_basic/_two_num_tool.py

@@ -0,0 +1,24 @@
+# -*- coding: utf-8 -*-
+"""
+Sample math tool for testing purposes.
+"""
+
+from massgen.tool._result import ExecutionResult, TextContent
+
+
+async def two_num_tool(x: int, y: int) -> ExecutionResult:
+    """Add two numbers together.
+
+    Args:
+        x: First number
+        y: Second number
+
+    Returns:
+        Sum of the two numbers
+    """
+    result = x + y
+    return ExecutionResult(
+        output_blocks=[
+            TextContent(data=f"The sum of {x} and {y} is {result}"),
+        ],
+    )

massgen/tool/_code_executors/__init__.py

@@ -0,0 +1,10 @@
+# -*- coding: utf-8 -*-
+"""Code execution tools."""
+
+from ._python_executor import run_python_script
+from ._shell_executor import run_shell_script
+
+__all__ = [
+    "run_python_script",
+    "run_shell_script",
+]

massgen/tool/_code_executors/_python_executor.py

@@ -0,0 +1,74 @@
+# -*- coding: utf-8 -*-
+"""Python code execution tool implementation."""
+
+import asyncio
+import os
+import sys
+import tempfile
+import uuid
+from typing import Any
+
+from .._result import ExecutionResult, TextContent
+
+
+async def run_python_script(
+    source_code: str,
+    max_duration: float = 300,
+    **extra_kwargs: Any,
+) -> ExecutionResult:
+    """Execute Python code in an isolated temporary environment.
+
+    The code runs in a temporary file that is cleaned up after execution.
+    Results must be printed to be captured in the output.
+
+    Args:
+        source_code: Python code to execute
+        max_duration: Maximum execution time in seconds (default: 300)
+
+    Returns:
+        ExecutionResult containing execution status and output
+    """
+
+    with tempfile.TemporaryDirectory() as temp_workspace:
+        script_file = os.path.join(temp_workspace, f"script_{uuid.uuid4().hex}.py")
+        with open(script_file, "w", encoding="utf-8") as file:
+            file.write(source_code)
+
+        process = await asyncio.create_subprocess_exec(
+            sys.executable,
+            "-u",
+            script_file,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+        )
+
+        try:
+            await asyncio.wait_for(process.wait(), timeout=max_duration)
+            stdout_data, stderr_data = await process.communicate()
+            stdout_text = stdout_data.decode("utf-8")
+            stderr_text = stderr_data.decode("utf-8")
+            exit_code = process.returncode
+
+        except asyncio.TimeoutError:
+            timeout_msg = f"TimeoutError: Execution exceeded {max_duration} seconds limit"
+            exit_code = -1
+            try:
+                process.terminate()
+                stdout_data, stderr_data = await process.communicate()
+                stdout_text = stdout_data.decode("utf-8")
+                stderr_text = stderr_data.decode("utf-8")
+                if stderr_text:
+                    stderr_text += f"\n{timeout_msg}"
+                else:
+                    stderr_text = timeout_msg
+            except ProcessLookupError:
+                stdout_text = ""
+                stderr_text = timeout_msg
+
+        return ExecutionResult(
+            output_blocks=[
+                TextContent(
+                    data=(f"<exit_code>{exit_code}</exit_code>" f"<stdout>{stdout_text}</stdout>" f"<stderr>{stderr_text}</stderr>"),
+                ),
+            ],
+        )

massgen/tool/_code_executors/_shell_executor.py

@@ -0,0 +1,61 @@
+# -*- coding: utf-8 -*-
+"""Shell command execution tool implementation."""
+
+import asyncio
+from typing import Any
+
+from .._result import ExecutionResult, TextContent
+
+
+async def run_shell_script(
+    shell_command: str,
+    max_duration: int = 300,
+    **extra_kwargs: Any,
+) -> ExecutionResult:
+    """Execute shell commands and capture output.
+
+    Args:
+        shell_command: Shell command to execute
+        max_duration: Maximum execution time in seconds (default: 300)
+
+    Returns:
+        ExecutionResult with exit code, stdout, and stderr
+    """
+
+    process = await asyncio.create_subprocess_shell(
+        shell_command,
+        stdout=asyncio.subprocess.PIPE,
+        stderr=asyncio.subprocess.PIPE,
+        bufsize=0,
+    )
+
+    try:
+        await asyncio.wait_for(process.wait(), timeout=max_duration)
+        stdout_data, stderr_data = await process.communicate()
+        stdout_text = stdout_data.decode("utf-8")
+        stderr_text = stderr_data.decode("utf-8")
+        exit_code = process.returncode
+
+    except asyncio.TimeoutError:
+        timeout_msg = f"TimeoutError: Command execution exceeded {max_duration} seconds limit"
+        exit_code = -1
+        try:
+            process.terminate()
+            stdout_data, stderr_data = await process.communicate()
+            stdout_text = stdout_data.decode("utf-8")
+            stderr_text = stderr_data.decode("utf-8")
+            if stderr_text:
+                stderr_text += f"\n{timeout_msg}"
+            else:
+                stderr_text = timeout_msg
+        except ProcessLookupError:
+            stdout_text = ""
+            stderr_text = timeout_msg
+
+    return ExecutionResult(
+        output_blocks=[
+            TextContent(
+                data=(f"<exit_code>{exit_code}</exit_code>" f"<stdout>{stdout_text}</stdout>" f"<stderr>{stderr_text}</stderr>"),
+            ),
+        ],
+    )

massgen/tool/_exceptions.py

@@ -0,0 +1,39 @@
+# -*- coding: utf-8 -*-
+"""Exception classes for tool operations."""
+
+
+class ToolException(Exception):
+    """Base exception for tool-related errors."""
+
+
+class InvalidToolArgumentsException(ToolException):
+    """Raised when tool receives invalid arguments."""
+
+    def __init__(self, error_msg: str):
+        self.error_msg = error_msg
+        super().__init__(self.error_msg)
+
+
+class ToolNotFoundException(ToolException):
+    """Raised when requested tool is not found."""
+
+    def __init__(self, tool_name: str):
+        self.tool_name = tool_name
+        super().__init__(f"Tool '{tool_name}' not found in registry")
+
+
+class ToolExecutionException(ToolException):
+    """Raised when tool execution fails."""
+
+    def __init__(self, tool_name: str, error_details: str):
+        self.tool_name = tool_name
+        self.error_details = error_details
+        super().__init__(f"Tool '{tool_name}' execution failed: {error_details}")
+
+
+class CategoryNotFoundException(ToolException):
+    """Raised when tool category is not found."""
+
+    def __init__(self, category_name: str):
+        self.category_name = category_name
+        super().__init__(f"Category '{category_name}' not found")

massgen/tool/_file_handlers/__init__.py

@@ -0,0 +1,10 @@
+# -*- coding: utf-8 -*-
+"""File handling tools."""
+
+from ._file_operations import append_file_content, read_file_content, save_file_content
+
+__all__ = [
+    "read_file_content",
+    "save_file_content",
+    "append_file_content",
+]

massgen/tool/_file_handlers/_file_operations.py

@@ -0,0 +1,218 @@
+# -*- coding: utf-8 -*-
+"""File operation tools for reading and writing."""
+
+import os
+from typing import List, Optional
+
+from .._result import ExecutionResult, TextContent
+
+
+async def read_file_content(
+    target_path: str,
+    line_range: Optional[List[int]] = None,
+) -> ExecutionResult:
+    """Read file contents with optional line range specification.
+
+    Args:
+        target_path: Path to the file to read
+        line_range: Optional [start, end] line numbers (1-based, inclusive).
+                   Use negative numbers for lines from end (e.g., [-100, -1])
+
+    Returns:
+        ExecutionResult with file content or error message
+    """
+    if not os.path.exists(target_path):
+        return ExecutionResult(
+            output_blocks=[
+                TextContent(
+                    data=f"Error: File '{target_path}' does not exist.",
+                ),
+            ],
+        )
+
+    if not os.path.isfile(target_path):
+        return ExecutionResult(
+            output_blocks=[
+                TextContent(
+                    data=f"Error: Path '{target_path}' is not a file.",
+                ),
+            ],
+        )
+
+    try:
+        with open(target_path, "r", encoding="utf-8") as file:
+            all_lines = file.readlines()
+
+        if line_range:
+            start_line, end_line = line_range
+
+            # Handle negative indices
+            if start_line < 0:
+                start_line = len(all_lines) + start_line + 1
+            if end_line < 0:
+                end_line = len(all_lines) + end_line + 1
+
+            # Validate range
+            if start_line < 1 or end_line > len(all_lines) or start_line > end_line:
+                return ExecutionResult(
+                    output_blocks=[
+                        TextContent(
+                            data=f"Error: Invalid line range {line_range} for file with {len(all_lines)} lines.",
+                        ),
+                    ],
+                )
+
+            # Extract lines (convert to 0-based indexing)
+            selected_lines = all_lines[start_line - 1 : end_line]
+            content_with_numbers = "".join(f"{i + start_line:4d}│ {line}" for i, line in enumerate(selected_lines))
+
+            return ExecutionResult(
+                output_blocks=[
+                    TextContent(
+                        data=f"Content of {target_path} (lines {start_line}-{end_line}):\n```\n{content_with_numbers}```",
+                    ),
+                ],
+            )
+        else:
+            full_content = "".join(all_lines)
+            return ExecutionResult(
+                output_blocks=[
+                    TextContent(
+                        data=f"Content of {target_path}:\n```\n{full_content}```",
+                    ),
+                ],
+            )
+
+    except Exception as error:
+        return ExecutionResult(
+            output_blocks=[
+                TextContent(
+                    data=f"Error reading file: {error}",
+                ),
+            ],
+        )
+
+
+async def save_file_content(
+    target_path: str,
+    file_content: str,
+    create_dirs: bool = True,
+) -> ExecutionResult:
+    """Write content to a file, creating directories if needed.
+
+    Args:
+        target_path: Path where file will be saved
+        file_content: Content to write to the file
+        create_dirs: Whether to create parent directories if they don't exist
+
+    Returns:
+        ExecutionResult indicating success or failure
+    """
+    try:
+        # Create parent directories if requested
+        if create_dirs:
+            parent_dir = os.path.dirname(target_path)
+            if parent_dir:
+                os.makedirs(parent_dir, exist_ok=True)
+
+        with open(target_path, "w", encoding="utf-8") as file:
+            file.write(file_content)
+
+        return ExecutionResult(
+            output_blocks=[
+                TextContent(
+                    data=f"Successfully wrote {len(file_content)} characters to {target_path}",
+                ),
+            ],
+        )
+
+    except Exception as error:
+        return ExecutionResult(
+            output_blocks=[
+                TextContent(
+                    data=f"Error writing file: {error}",
+                ),
+            ],
+        )
+
+
+async def append_file_content(
+    target_path: str,
+    additional_content: str,
+    line_position: Optional[int] = None,
+) -> ExecutionResult:
+    """Append content to a file or insert at specific line.
+
+    Args:
+        target_path: Path to the file
+        additional_content: Content to append or insert
+        line_position: Optional line number to insert at (1-based).
+                      If None, appends to end of file.
+
+    Returns:
+        ExecutionResult indicating success or failure
+    """
+    if not os.path.exists(target_path):
+        return ExecutionResult(
+            output_blocks=[
+                TextContent(
+                    data=f"Error: File '{target_path}' does not exist.",
+                ),
+            ],
+        )
+
+    try:
+        with open(target_path, "r", encoding="utf-8") as file:
+            current_lines = file.readlines()
+
+        if line_position is None:
+            # Append to end
+            with open(target_path, "a", encoding="utf-8") as file:
+                file.write(additional_content)
+
+            return ExecutionResult(
+                output_blocks=[
+                    TextContent(
+                        data=f"Successfully appended {len(additional_content)} characters to {target_path}",
+                    ),
+                ],
+            )
+        else:
+            # Insert at specific line
+            if line_position < 1 or line_position > len(current_lines) + 1:
+                return ExecutionResult(
+                    output_blocks=[
+                        TextContent(
+                            data=f"Error: Invalid line position {line_position} for file with {len(current_lines)} lines.",
+                        ),
+                    ],
+                )
+
+            # Insert content (convert to 0-based index)
+            insert_idx = line_position - 1
+
+            # Ensure content ends with newline for proper insertion
+            if not additional_content.endswith("\n"):
+                additional_content += "\n"
+
+            current_lines.insert(insert_idx, additional_content)
+
+            with open(target_path, "w", encoding="utf-8") as file:
+                file.writelines(current_lines)
+
+            return ExecutionResult(
+                output_blocks=[
+                    TextContent(
+                        data=f"Successfully inserted content at line {line_position} in {target_path}",
+                    ),
+                ],
+            )
+
+    except Exception as error:
+        return ExecutionResult(
+            output_blocks=[
+                TextContent(
+                    data=f"Error modifying file: {error}",
+                ),
+            ],
+        )