massgen 0.0.3__py3-none-any.whl

massgen/v1/tools.py

@@ -0,0 +1,160 @@
+import inspect
+import json
+import random
+import subprocess
+import sys
+import time
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Any, Union, Optional, Dict, List
+import ast
+import operator
+import math
+
+# Global tool registry
+register_tool = {}
+
+# Mock functions removed - actual functionality is implemented in agent classes
+
+
+def python_interpreter(code: str, timeout: Optional[int] = 10) -> Dict[str, Any]:
+    """
+    Execute Python code in an isolated subprocess and return its output.
+
+    Args:
+        code: The Python code string to execute
+        timeout: Maximum execution time in seconds (default: 10, Must be less than 60 seconds)
+
+    Returns:
+        A dictionary containing:
+        - 'stdout': Standard output from the code execution
+        - 'stderr': Standard error from the code execution
+        - 'returncode': Exit code of the process (0 for success)
+        - 'success': Boolean indicating if execution was successful
+        - 'error': Error message if execution failed
+    """
+    # Ensure timeout is between 0 and 60 seconds
+    timeout = max(min(timeout, 60), 0)
+    try:
+        # Run the code in a separate Python process
+        result = subprocess.run(
+            [sys.executable, "-c", code],
+            capture_output=True,
+            text=True,
+            timeout=timeout,
+        )
+
+        return json.dumps(
+            {
+                "stdout": result.stdout,
+                "stderr": result.stderr,
+                "returncode": result.returncode,
+                "success": result.returncode == 0,
+                "error": None,
+            }
+        )
+
+    except subprocess.TimeoutExpired:
+        return json.dumps(
+            {
+                "stdout": "",
+                "stderr": "",
+                "returncode": -1,
+                "success": False,
+                "error": f"Code execution timed out after {timeout} seconds",
+            }
+        )
+
+    except Exception as e:
+        return json.dumps(
+            {
+                "stdout": "",
+                "stderr": "",
+                "returncode": -1,
+                "success": False,
+                "error": f"Failed to execute code: {str(e)}",
+            }
+        )
+
+
+def calculator(expression: str) -> float:
+    """
+    Mathematical expression to evaluate (e.g., '2 + 3 * 4', 'sqrt(16)', 'sin(pi/2)')
+    """
+    safe_operators = {
+        ast.Add: operator.add,
+        ast.Sub: operator.sub,
+        ast.Mult: operator.mul,
+        ast.Div: operator.truediv,
+        ast.Pow: operator.pow,
+        ast.USub: operator.neg,
+        ast.UAdd: operator.pos,
+        ast.Mod: operator.mod,
+    }
+
+    # Safe functions
+    safe_functions = {
+        "abs": abs,
+        "round": round,
+        "max": max,
+        "min": min,
+        "sum": sum,
+        "sqrt": math.sqrt,
+        "sin": math.sin,
+        "cos": math.cos,
+        "tan": math.tan,
+        "log": math.log,
+        "log10": math.log10,
+        "exp": math.exp,
+        "pi": math.pi,
+        "e": math.e,
+    }
+
+    def _safe_eval(node):
+        """Safely evaluate an AST node"""
+        if isinstance(node, ast.Constant):  # Numbers
+            return node.value
+        elif isinstance(node, ast.Name):  # Variables/constants
+            if node.id in safe_functions:
+                return safe_functions[node.id]
+            else:
+                raise ValueError(f"Unknown variable: {node.id}")
+        elif isinstance(node, ast.BinOp):  # Binary operations
+            left = _safe_eval(node.left)
+            right = _safe_eval(node.right)
+            if type(node.op) in safe_operators:
+                return safe_operators[type(node.op)](left, right)
+            else:
+                raise ValueError(f"Unsupported operation: {type(node.op)}")
+        elif isinstance(node, ast.UnaryOp):  # Unary operations
+            operand = _safe_eval(node.operand)
+            if type(node.op) in safe_operators:
+                return safe_operators[type(node.op)](operand)
+            else:
+                raise ValueError(f"Unsupported unary operation: {type(node.op)}")
+        elif isinstance(node, ast.Call):  # Function calls
+            func = _safe_eval(node.func)
+            args = [_safe_eval(arg) for arg in node.args]
+            return func(*args)
+        else:
+            raise ValueError(f"Unsupported node type: {type(node)}")
+
+    try:
+        # Parse the expression
+        tree = ast.parse(expression, mode="eval")
+
+        # Evaluate safely
+        result = _safe_eval(tree.body)
+
+        return {"expression": expression, "result": result, "success": True}
+
+    except Exception as e:
+        return {"expression": expression, "error": str(e), "success": False}
+
+
+# Register tools in the global registry
+register_tool["python_interpreter"] = python_interpreter
+register_tool["calculator"] = calculator
+
+if __name__ == "__main__":
+    print(calculator("24423 + 312 * log(10)"))

massgen/v1/types.py

@@ -0,0 +1,245 @@
+"""
+MassGen System Types
+
+This module contains all the core type definitions and dataclasses
+used throughout the MassGen framework.
+"""
+
+import time
+from dataclasses import dataclass, field, asdict
+from typing import Any, Dict, List, Optional
+from abc import ABC, abstractmethod
+
+
+@dataclass
+class AnswerRecord:
+    """Represents a single answer record in an agent's update history."""
+
+    timestamp: float
+    answer: str
+    status: str
+
+    def __post_init__(self):
+        """Ensure timestamp is set if not provided."""
+        if not self.timestamp:
+            self.timestamp = time.time()
+
+
+@dataclass
+class VoteRecord:
+    """Records a vote cast by an agent."""
+
+    voter_id: int
+    target_id: int
+    reason: str = ""  # the full response text that led to this vote
+    timestamp: float = 0.0
+
+    def __post_init__(self):
+        """Ensure timestamp is set if not provided."""
+        if not self.timestamp:
+            import time
+
+            self.timestamp = time.time()
+
+
+@dataclass
+class ModelConfig:
+    """Configuration for agent model parameters."""
+
+    model: Optional[str] = None
+    tools: Optional[List[str]] = None
+    max_retries: int = 10  # max retries for each LLM call
+    max_rounds: int = 10  # max round for task
+    max_tokens: Optional[int] = None
+    temperature: Optional[float] = None
+    top_p: Optional[float] = None
+    inference_timeout: Optional[float] = 180  # seconds
+    stream: bool = True  # whether to stream the response
+
+
+@dataclass
+class TaskInput:
+    """Represents a task to be processed by the MassGen system."""
+
+    question: str
+    context: Dict[str, Any] = field(
+        default_factory=dict
+    )  # may support more information in the future, like images
+    task_id: Optional[str] = None
+
+
+@dataclass
+class SystemState:
+    """Overall state of the MassGen orchestrator.
+    Simplified phases:
+    - collaboration: agents are working together to solve the task
+    - completed: the representative agent has presented the solution and the task is completed
+    """
+
+    task: Optional[TaskInput] = None
+    phase: str = "collaboration"  # "collaboration", "debate", "completed"
+    start_time: Optional[float] = None
+    end_time: Optional[float] = None
+    consensus_reached: bool = False
+    representative_agent_id: Optional[int] = None
+
+
+@dataclass
+class AgentState:
+    """Represents the current state of an agent in the MassGen system."""
+
+    agent_id: int
+    status: str = "working"  # "working", "voted", "failed"
+    curr_answer: str = ""  # the latest answer of the agent's work
+    updated_answers: List[AnswerRecord] = field(
+        default_factory=list
+    )  # a list of answer records
+    curr_vote: Optional[VoteRecord] = (
+        None  # Which agent's solution this agent voted for
+    )
+    cast_votes: List[VoteRecord] = field(default_factory=list)  # a list of vote records
+    seen_updates_timestamps: Dict[int, float] = field(
+        default_factory=dict
+    )  # agent_id -> last_seen_timestamp
+    chat_history: List[Dict[str, Any]] = field(
+        default_factory=list
+    )  # a list of conversation records
+    chat_round: int = 0  # the number of chat rounds the agent has participated in
+    execution_start_time: Optional[float] = None
+    execution_end_time: Optional[float] = None
+
+    @property
+    def execution_time(self) -> Optional[float]:
+        """Calculate execution time if both start and end times are available."""
+        if self.execution_start_time and self.execution_end_time:
+            return self.execution_end_time - self.execution_start_time
+        return None
+
+    def add_update(self, answer: str, timestamp: Optional[float] = None):
+        """Add an update to the agent's history."""
+        if timestamp is None:
+            timestamp = time.time()
+
+        record = AnswerRecord(
+            timestamp=timestamp,
+            answer=answer,
+            status=self.status,
+        )
+        self.updated_answers.append(record)
+        self.curr_answer = answer
+
+    def mark_updates_seen(self, agent_updates: Dict[int, float]):
+        """Mark updates from other agents as seen."""
+        for agent_id, timestamp in agent_updates.items():
+            if agent_id != self.agent_id:  # Don't track own updates
+                self.seen_updates_timestamps[agent_id] = timestamp
+
+    def has_unseen_updates(self, other_agent_updates: Dict[int, float]) -> bool:
+        """Check if there are unseen updates from other agents."""
+        for agent_id, timestamp in other_agent_updates.items():
+            if agent_id != self.agent_id:  # Don't check own updates
+                last_seen = self.seen_updates_timestamps.get(agent_id, 0)
+                if timestamp > last_seen:
+                    return True
+        return False
+
+
+@dataclass
+class AgentResponse:
+    """Response from an agent's process_message function."""
+
+    text: str
+    code: List[str] = field(default_factory=list)
+    citations: List[Dict[str, Any]] = field(default_factory=list)
+    function_calls: List[Dict[str, Any]] = field(default_factory=list)
+
+
+@dataclass
+class LogEntry:
+    """Represents a single log entry in the MassGen system."""
+
+    timestamp: float
+    event_type: str  # e.g., "agent_answer_update", "voting", "phase_change", etc.
+    agent_id: Optional[int]
+    phase: str
+    data: Dict[str, Any]
+    session_id: Optional[str] = None
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary for JSON serialization."""
+        return asdict(self)
+
+
+@dataclass
+class StreamingDisplayConfig:
+    """Configuration for streaming display system."""
+
+    display_enabled: bool = True
+    max_lines: int = 10
+    save_logs: bool = True
+    stream_callback: Optional[Any] = None  # Callable, but avoid circular imports
+
+
+@dataclass
+class LoggingConfig:
+    """Configuration for logging system."""
+
+    log_dir: str = "logs"
+    session_id: Optional[str] = None
+    non_blocking: bool = False
+
+
+@dataclass
+class OrchestratorConfig:
+    """Configuration for MassGen orchestrator."""
+
+    max_duration: int = 600
+    consensus_threshold: float = 0.0
+    max_debate_rounds: int = 1
+    status_check_interval: float = 2.0
+    thread_pool_timeout: int = 5
+
+
+@dataclass
+class AgentConfig:
+    """Complete configuration for a single agent."""
+
+    agent_id: int
+    agent_type: str  # "openai", "gemini", "grok"
+    model_config: ModelConfig
+
+    def __post_init__(self):
+        """Validate agent configuration."""
+        if self.agent_type not in ["openai", "gemini", "grok"]:
+            raise ValueError(
+                f"Invalid agent_type: {self.agent_type}. Must be one of: openai, gemini, grok"
+            )
+
+
+@dataclass
+class MassConfig:
+    """Complete MassGen system configuration."""
+
+    orchestrator: OrchestratorConfig = field(default_factory=OrchestratorConfig)
+    agents: List[AgentConfig] = field(default_factory=list)
+    streaming_display: StreamingDisplayConfig = field(
+        default_factory=StreamingDisplayConfig
+    )
+    logging: LoggingConfig = field(default_factory=LoggingConfig)
+    task: Optional[Dict[str, Any]] = None  # Task-specific configuration
+
+    def validate(self) -> bool:
+        """Validate the complete configuration."""
+        if not self.agents:
+            raise ValueError("At least one agent must be configured")
+
+        # Check for duplicate agent IDs
+        agent_ids = [agent.agent_id for agent in self.agents]
+        if len(agent_ids) != len(set(agent_ids)):
+            raise ValueError("Agent IDs must be unique")
+
+        # Validate consensus threshold
+        if not 0.0 <= self.orchestrator.consensus_threshold <= 1.0:
+            raise ValueError("Consensus threshold must be between 0.0 and 1.0")
+
+        return True

massgen/v1/utils.py

@@ -0,0 +1,199 @@
+import inspect
+import json
+import random
+import subprocess
+import sys
+import time
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Any, Union, Optional, Dict, List
+import ast
+import operator
+import math
+
+# Model mappings and constants
+MODEL_MAPPINGS = {
+    "openai": [
+        # GPT-4.1 variants
+        "gpt-4.1",
+        "gpt-4.1-mini",
+        # GPT-4o variants
+        "gpt-4o-mini",
+        "gpt-4o",
+        # o1
+        "o1",  # -> o1-2024-12-17
+        # o3
+        "o3",
+        "o3-low",
+        "o3-medium",
+        "o3-high",
+        # o3 mini
+        "o3-mini",
+        "o3-mini-low",
+        "o3-mini-medium",
+        "o3-mini-high",
+        # o4 mini
+        "o4-mini",
+        "o4-mini-low",
+        "o4-mini-medium",
+        "o4-mini-high",
+    ],
+    "gemini": [
+        "gemini-2.5-flash",
+        "gemini-2.5-pro",
+    ],
+    "grok": [
+        "grok-3-mini",
+        "grok-3",
+        "grok-4",
+    ],
+}
+
+
+def get_agent_type_from_model(model: str) -> str:
+    """
+    Determine the agent type based on the model name.
+
+    Args:
+        model: The model name (e.g., "gpt-4", "gemini-pro", "grok-1")
+
+    Returns:
+        Agent type string ("openai", "gemini", "grok")
+    """
+    if not model:
+        return "openai"  # Default to OpenAI
+
+    model_lower = model.lower()
+
+    for key, models in MODEL_MAPPINGS.items():
+        if model_lower in models:
+            return key
+    raise ValueError(f"Unknown model: {model}")
+
+
+def get_available_models() -> list:
+    """Get a flat list of all available model names."""
+    all_models = []
+    for models in MODEL_MAPPINGS.values():
+        all_models.extend(models)
+    return all_models
+
+
+def generate_random_id(length: int = 24) -> str:
+    """Generate a random ID string."""
+    characters = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789"
+    return "".join(random.choice(characters) for _ in range(length))
+
+
+# Utility functions (originally from util.py)
+def execute_function_calls(function_calls, tool_mapping):
+    """Execute function calls and return formatted outputs for the conversation."""
+    function_outputs = []
+    for function_call in function_calls:
+        try:
+            # Get the function from tool mapping
+            target_function = None
+            function_name = function_call.get("name")
+
+            # Look up function in tool_mapping
+            if function_name in tool_mapping:
+                target_function = tool_mapping[function_name]
+            else:
+                # Handle error case
+                error_output = {
+                    "type": "function_call_output",
+                    "call_id": function_call.get("call_id"),
+                    "output": f"Error: Function '{function_name}' not found in tool mapping",
+                }
+                function_outputs.append(error_output)
+                continue
+
+            # Parse arguments and execute function
+            if isinstance(function_call.get("arguments", {}), str):
+                arguments = json.loads(function_call.get("arguments", "{}"))
+            elif isinstance(function_call.get("arguments", {}), dict):
+                arguments = function_call.get("arguments", {})
+            else:
+                raise ValueError(
+                    f"Unknown arguments type: {type(function_call.get('arguments', {}))}"
+                )
+            result = target_function(**arguments)
+
+            # Format the output according to Responses API requirements
+            function_output = {
+                "type": "function_call_output",
+                "call_id": function_call.get("call_id"),
+                "output": str(result),
+            }
+            function_outputs.append(function_output)
+
+            # print(f"Executed function: {function_name}({arguments}) -> {result}")
+
+        except Exception as e:
+            # Handle execution errors
+            error_output = {
+                "type": "function_call_output",
+                "call_id": function_call.get("call_id"),
+                "output": f"Error executing function: {str(e)}",
+            }
+            function_outputs.append(error_output)
+            # print(f"Error executing function {function_name}: {e}")
+
+    return function_outputs
+
+
+def function_to_json(func) -> dict:
+    """
+    Converts a Python function into a JSON-serializable dictionary
+    that describes the function's signature, including its name,
+    description, and parameters.
+
+    Args:
+        func: The function to be converted.
+
+    Returns:
+        A dictionary representing the function's signature in JSON format.
+    """
+    type_map = {
+        str: "string",
+        int: "integer",
+        float: "number",
+        bool: "boolean",
+        list: "array",
+        dict: "object",
+        type(None): "null",
+    }
+
+    try:
+        signature = inspect.signature(func)
+    except ValueError as e:
+        raise ValueError(
+            f"Failed to get signature for function {func.__name__}: {str(e)}"
+        )
+
+    parameters = {}
+    for param in signature.parameters.values():
+        try:
+            param_type = type_map.get(param.annotation, "string")
+        except KeyError as e:
+            raise KeyError(
+                f"Unknown type annotation {param.annotation} for parameter {param.name}: {str(e)}"
+            )
+        parameters[param.name] = {"type": param_type}
+
+    required = [
+        param.name
+        for param in signature.parameters.values()
+        if param.default == inspect._empty
+    ]
+
+    return {
+        "type": "function",
+        "name": func.__name__,
+        "description": func.__doc__ or "",
+        "parameters": {
+            "type": "object",
+            "properties": parameters,
+            "required": required,
+        },
+    }