massgen 0.0.3__py3-none-any.whl → 0.1.0__py3-none-any.whl

massgen/adapters/tests/test_agent_adapter.py

@@ -0,0 +1,128 @@
+# -*- coding: utf-8 -*-
+"""
+Unit tests for base AgentAdapter class.
+"""
+import pytest
+
+from massgen.adapters.base import AgentAdapter
+
+
+class MockAdapter(AgentAdapter):
+    """Mock adapter for testing."""
+
+    async def execute_streaming(self, messages, tools, **kwargs):
+        """Mock implementation."""
+        content = "Mock response"
+        async for chunk in self.simulate_streaming(content):
+            yield chunk
+
+
+@pytest.mark.asyncio
+async def test_simulate_streaming():
+    """Test streaming simulation."""
+    adapter = MockAdapter()
+
+    content = "Hello world"
+    chunks = []
+
+    async for chunk in adapter.simulate_streaming(content):
+        chunks.append(chunk)
+
+    # Should have content chunks + complete_message + done
+    assert len(chunks) > 0
+    assert any(c.type == "content" for c in chunks)
+    assert any(c.type == "complete_message" for c in chunks)
+    assert any(c.type == "done" for c in chunks)
+
+    # Complete message should have content
+    complete_chunks = [c for c in chunks if c.type == "complete_message"]
+    assert len(complete_chunks) == 1
+    assert complete_chunks[0].complete_message["content"] == content
+
+
+@pytest.mark.asyncio
+async def test_simulate_streaming_with_tool_calls():
+    """Test streaming simulation with tool calls."""
+    adapter = MockAdapter()
+
+    content = "Using tools"
+    tool_calls = [
+        {
+            "id": "call_1",
+            "function": {"name": "search", "arguments": '{"query": "test"}'},
+        },
+    ]
+
+    chunks = []
+    async for chunk in adapter.simulate_streaming(content, tool_calls):
+        chunks.append(chunk)
+
+    # Should have tool_calls chunk
+    tool_chunks = [c for c in chunks if c.type == "tool_calls"]
+    assert len(tool_chunks) == 1
+    assert tool_chunks[0].tool_calls == tool_calls
+
+    # Complete message should include tool calls
+    complete_chunks = [c for c in chunks if c.type == "complete_message"]
+    assert complete_chunks[0].complete_message["tool_calls"] == tool_calls
+
+
+def test_convert_messages_default():
+    """Test default message conversion (passthrough)."""
+    adapter = MockAdapter()
+
+    messages = [
+        {"role": "user", "content": "Hello"},
+        {"role": "assistant", "content": "Hi"},
+    ]
+
+    result = adapter.convert_messages_from_massgen(messages)
+    assert result == messages
+
+
+def test_convert_tools_default():
+    """Test default tool conversion (passthrough)."""
+    adapter = MockAdapter()
+
+    tools = [
+        {
+            "type": "function",
+            "function": {"name": "search", "description": "Search tool"},
+        },
+    ]
+
+    result = adapter.convert_tools_from_massgen(tools)
+    assert result == tools
+
+
+def test_is_stateful_default():
+    """Test default stateful behavior."""
+    adapter = MockAdapter()
+    assert adapter.is_stateful() is False
+
+
+def test_clear_history():
+    """Test clearing conversation history."""
+    adapter = MockAdapter()
+
+    # Add some history
+    adapter._conversation_history = [{"role": "user", "content": "test"}]
+
+    # Clear it
+    adapter.clear_history()
+
+    assert len(adapter._conversation_history) == 0
+
+
+def test_reset_state():
+    """Test resetting adapter state."""
+    adapter = MockAdapter()
+
+    # Add some history
+    adapter._conversation_history = [{"role": "user", "content": "test"}]
+
+    # Reset
+    adapter.reset_state()
+
+    # Should clear history
+    assert len(adapter._conversation_history) == 0

massgen/adapters/utils/__init__.py

@@ -0,0 +1,2 @@
+# -*- coding: utf-8 -*-
+"""Utility functions for adapters."""

massgen/adapters/utils/ag2_utils.py

@@ -0,0 +1,236 @@
+# -*- coding: utf-8 -*-
+"""
+Utility functions for AG2 (AutoGen) adapter.
+"""
+import os
+import time
+
+# Suppress autogen deprecation warnings
+import warnings
+from typing import Any, Dict, List
+
+warnings.filterwarnings("ignore", category=DeprecationWarning, module="autogen")
+warnings.filterwarnings("ignore", message=".*jsonschema.*")
+warnings.filterwarnings("ignore", message=".*Pydantic.*")
+
+from autogen import AssistantAgent, ConversableAgent, LLMConfig  # noqa: E402
+
+
+def setup_api_keys() -> None:
+    """Set up API keys for AG2 compatibility."""
+    # Copy GEMINI_API_KEY to GOOGLE_GEMINI_API_KEY if it exists
+    if "GEMINI_API_KEY" in os.environ and "GOOGLE_GEMINI_API_KEY" not in os.environ:
+        os.environ["GOOGLE_GEMINI_API_KEY"] = os.environ["GEMINI_API_KEY"]
+
+
+def validate_agent_config(cfg: Dict[str, Any], require_llm_config: bool = True) -> None:
+    """
+    Validate required fields in agent configuration.
+
+    Args:
+        cfg: Agent configuration dict
+        require_llm_config: If True, llm_config is required. If False, it's optional.
+    """
+    if require_llm_config and "llm_config" not in cfg:
+        raise ValueError("Each AG2 agent configuration must include 'llm_config'.")
+
+    if "name" not in cfg:
+        raise ValueError("Each AG2 agent configuration must include 'name'.")
+
+
+def create_llm_config(llm_config_data: Any) -> LLMConfig:
+    """
+    Create LLMConfig from dict or list format.
+
+    Supports new AG2 syntax:
+    - Single dict: LLMConfig({'model': 'gpt-4', 'api_key': '...'})
+    - List of dicts: LLMConfig({'model': 'gpt-4', ...}, {'model': 'gpt-3.5', ...})
+    """
+    if isinstance(llm_config_data, list):
+        # YAML format: llm_config: [{...}, {...}]
+        return LLMConfig(*llm_config_data)
+    elif isinstance(llm_config_data, dict):
+        # YAML format: llm_config: {model: 'gpt-4o', ...}
+        return LLMConfig(llm_config_data)
+    else:
+        raise ValueError(f"llm_config must be a dict or list, got {type(llm_config_data)}")
+
+
+def create_code_executor(executor_config: Dict[str, Any]) -> Any:
+    """Create code executor from configuration."""
+    executor_type = executor_config.get("type")
+
+    if not executor_type:
+        raise ValueError("code_execution_config.executor must include 'type' field")
+
+    # Remove 'type' from config before passing to executor
+    executor_params = {k: v for k, v in executor_config.items() if k != "type"}
+
+    # Create appropriate executor based on type
+    if executor_type == "LocalCommandLineCodeExecutor":
+        from autogen.coding import LocalCommandLineCodeExecutor
+
+        return LocalCommandLineCodeExecutor(**executor_params)
+
+    elif executor_type == "DockerCommandLineCodeExecutor":
+        from autogen.coding import DockerCommandLineCodeExecutor
+
+        return DockerCommandLineCodeExecutor(**executor_params)
+
+    elif executor_type == "YepCodeCodeExecutor":
+        from autogen.coding import YepCodeCodeExecutor
+
+        return YepCodeCodeExecutor(**executor_params)
+
+    elif executor_type == "JupyterCodeExecutor":
+        from autogen.coding.jupyter import JupyterCodeExecutor
+
+        return JupyterCodeExecutor(**executor_params)
+
+    else:
+        raise ValueError(
+            f"Unsupported code executor type: {executor_type}. " f"Supported types: LocalCommandLineCodeExecutor, DockerCommandLineCodeExecutor, " f"YepCodeCodeExecutor, JupyterCodeExecutor",
+        )
+
+
+def build_agent_kwargs(cfg: Dict[str, Any], llm_config: LLMConfig, code_executor: Any = None) -> Dict[str, Any]:
+    """Build kwargs for agent initialization."""
+    agent_kwargs = {
+        "name": cfg["name"],
+        "system_message": cfg.get("system_message", "You are a helpful AI assistant."),
+        "human_input_mode": "NEVER",
+        "llm_config": llm_config,
+    }
+
+    if code_executor is not None:
+        agent_kwargs["code_execution_config"] = {"executor": code_executor}
+
+    return agent_kwargs
+
+
+def setup_agent_from_config(config: Dict[str, Any], default_llm_config: Any = None) -> ConversableAgent:
+    """
+    Set up a ConversableAgent from configuration.
+
+    Args:
+        config: Agent configuration dict
+        default_llm_config: Default llm_config to use if agent doesn't provide one
+
+    Returns:
+        ConversableAgent or AssistantAgent instance
+    """
+    cfg = config.copy()
+
+    # Check if llm_config is provided in agent config
+    has_llm_config = "llm_config" in cfg
+
+    # Validate configuration (llm_config optional if default provided)
+    validate_agent_config(cfg, require_llm_config=not default_llm_config)
+
+    # Extract agent type
+    agent_type = cfg.pop("type", "conversable")
+
+    # Create LLM config
+    if has_llm_config:
+        llm_config = create_llm_config(cfg.pop("llm_config"))
+    elif default_llm_config:
+        llm_config = create_llm_config(default_llm_config)
+    else:
+        raise ValueError("No llm_config provided for agent and no default_llm_config available")
+
+    # Create code executor if configured
+    code_executor = None
+    if "code_execution_config" in cfg:
+        code_exec_config = cfg.pop("code_execution_config")
+        if "executor" in code_exec_config:
+            code_executor = create_code_executor(code_exec_config["executor"])
+
+    # Build agent kwargs
+    agent_kwargs = build_agent_kwargs(cfg, llm_config, code_executor)
+
+    # Create appropriate agent
+    if agent_type == "assistant":
+        return AssistantAgent(**agent_kwargs)
+    elif agent_type == "conversable":
+        return ConversableAgent(**agent_kwargs)
+    else:
+        raise ValueError(
+            f"Unsupported AG2 agent type: {agent_type}. Use 'assistant' or 'conversable' for ag2 agents.",
+        )
+
+
+def get_group_initial_message() -> Dict[str, Any] | None:
+    """
+    Create the initial system message for group chat.
+
+    Returns:
+        Dict with role and content for initial system message
+    """
+    initial_message = f"""
+    CURRENT ANSWER from multiple agents for final response to a message is given.
+    Different agents may have different builtin tools and capabilities.
+    Does the best CURRENT ANSWER address the ORIGINAL MESSAGE well?
+
+    If CURRENT ANSWER is given, digest existing answers, combine their strengths, and do additional work to address their weaknesses.
+    if you think CURRENT ANSWER is good enough, you can also use it as your answer.
+
+    *Note*: The CURRENT TIME is **{time.strftime("%Y-%m-%d %H:%M:%S")}**.
+    """
+
+    # Not real system message. Can't send system message to group chat.
+    # When a non function/tool message is sent to an agent in group chat, it will be treated as user message.
+    return {"role": "system", "content": initial_message}
+
+
+def get_user_agent_tool_call_message() -> str:
+    system_message = """
+    You are the User agent overseeing a team of expert agents.
+    They worked together to create an improved answer to the ORIGINAL MESSAGE based on CURRENT ANSWER (if given).
+
+    Does CURRENT ANSWER address the ORIGINAL MESSAGE well? If YES, use the `vote` tool to record your vote and skip the `new_answer` tool.
+    Otherwise, find the final improved answer generated by team and use the `new_answer` tool to provide it as your final answer to the ORIGINAL MESSAGE.
+
+    When CURRENT ANSWER section is not available and a new answer is provided by the team of experts,
+    you should use the `new_answer` tool instead of `vote` tool.
+
+    You MUST ONLY use one of the two tools (`vote` or `new_answer`) ONCE to respond.
+    """
+
+    return system_message
+
+
+def get_user_agent_default_system_message() -> str:
+    system_message = """
+    "MUST say 'TERMINATE' when the original request is well answered. Do NOT do anything else."
+    """
+    return system_message
+
+
+def get_user_agent_default_description() -> str:
+    description = """
+    ALWAYS check if other agents still needs to be selected before selected this agent.
+    MUST ONLY be selected when the original request is well answered and the conversation should terminate.
+    """
+
+    return description
+
+
+def postprocess_group_chat_results(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    for message in messages:
+        if message["content"]:
+            message["content"] = f"<SENDER>: {message['name']} </SENDER> \n" + message["content"]
+        message["role"] = "assistant"
+
+    return messages
+
+
+def unregister_tools_for_agent(tools: List[Dict[str, Any]], agent: ConversableAgent) -> None:
+    """Unregister all tools from single agent."""
+    for tool in tools:
+        agent.update_tool_signature(tool_sig=tool, is_remove=True, silent_override=True)
+
+
+def register_tools_for_agent(tools: List[Dict[str, Any]], agent: ConversableAgent) -> None:
+    """Register all tools to single agent."""
+    for tool in tools:
+        agent.update_tool_signature(tool_sig=tool, is_remove=False, silent_override=True)

massgen/adapters/utils/tests/test_ag2_utils.py

@@ -0,0 +1,138 @@
+# -*- coding: utf-8 -*-
+"""
+Unit tests for AG2 utility functions.
+"""
+from unittest.mock import MagicMock
+
+from massgen.adapters.utils.ag2_utils import (
+    create_llm_config,
+    get_group_initial_message,
+    postprocess_group_chat_results,
+    register_tools_for_agent,
+    unregister_tools_for_agent,
+)
+
+
+def test_create_llm_config_from_dict():
+    """Test creating LLMConfig from dictionary."""
+    config_dict = {
+        "api_type": "openai",
+        "model": "gpt-4o",
+        "temperature": 0.7,
+    }
+
+    llm_config = create_llm_config(config_dict)
+
+    # Should create LLMConfig instance
+    assert llm_config is not None
+    # Should contain config_list
+    assert hasattr(llm_config, "config_list")
+
+
+def test_create_llm_config_from_list():
+    """Test creating LLMConfig from list of configs."""
+    config_list = [
+        {"api_type": "openai", "model": "gpt-4o"},
+        {"api_type": "google", "model": "gemini-pro"},
+    ]
+
+    llm_config = create_llm_config(config_list)
+
+    # Should create LLMConfig instance
+    assert llm_config is not None
+    assert hasattr(llm_config, "config_list")
+
+
+def test_postprocess_group_chat_results():
+    """Test postprocessing of group chat results."""
+    messages = [
+        {"name": "Agent1", "content": "Hello", "role": "user"},
+        {"name": "Agent2", "content": "Hi there", "role": "user"},
+    ]
+
+    result = postprocess_group_chat_results(messages)
+
+    # Should add sender tags to content
+    assert "<SENDER>: Agent1 </SENDER>" in result[0]["content"]
+    assert "<SENDER>: Agent2 </SENDER>" in result[1]["content"]
+
+    # Should change role to assistant
+    assert result[0]["role"] == "assistant"
+    assert result[1]["role"] == "assistant"
+
+
+def test_postprocess_group_chat_results_empty_content():
+    """Test postprocessing with empty content."""
+    messages = [
+        {"name": "Agent1", "content": "", "role": "user"},
+        {"name": "Agent2", "content": None, "role": "user"},
+    ]
+
+    result = postprocess_group_chat_results(messages)
+
+    # Should not add sender tags for empty content
+    assert result[0]["content"] == ""
+    assert result[1]["content"] is None
+
+    # Should still change role
+    assert result[0]["role"] == "assistant"
+    assert result[1]["role"] == "assistant"
+
+
+def test_get_group_initial_message():
+    """Test getting initial message for group chat."""
+    message = get_group_initial_message()
+
+    # Should return a dict with role and content
+    assert isinstance(message, dict)
+    assert "role" in message
+    assert "content" in message
+
+    # Should be system role
+    assert message["role"] == "system"
+
+    # Content should mention key concepts
+    assert "CURRENT ANSWER" in message["content"]
+    assert "ORIGINAL MESSAGE" in message["content"]
+
+
+def test_register_tools_for_agent():
+    """Test registering tools with agent."""
+    mock_agent = MagicMock()
+    tools = [
+        {
+            "type": "function",
+            "function": {"name": "search", "description": "Search tool"},
+        },
+        {
+            "type": "function",
+            "function": {"name": "calc", "description": "Calculator tool"},
+        },
+    ]
+
+    register_tools_for_agent(tools, mock_agent)
+
+    # Should call update_tool_signature for each tool
+    assert mock_agent.update_tool_signature.call_count == len(tools)
+
+    # Should be called with is_remove=False
+    for call in mock_agent.update_tool_signature.call_args_list:
+        assert call[1]["is_remove"] is False
+
+
+def test_unregister_tools_for_agent():
+    """Test unregistering tools from agent."""
+    mock_agent = MagicMock()
+    tools = [
+        {
+            "type": "function",
+            "function": {"name": "search", "description": "Search tool"},
+        },
+    ]
+
+    unregister_tools_for_agent(tools, mock_agent)
+
+    # Should call update_tool_signature with is_remove=True
+    mock_agent.update_tool_signature.assert_called_once()
+    call_kwargs = mock_agent.update_tool_signature.call_args[1]
+    assert call_kwargs["is_remove"] is True