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

massgen/tests/test_conversation_memory.py

@@ -0,0 +1,382 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Tests for ConversationMemory implementation.
+
+This module tests the in-memory conversation storage functionality,
+including adding, retrieving, deleting, and managing messages.
+"""
+
+import pytest
+
+from massgen.memory import ConversationMemory
+
+
+@pytest.mark.asyncio
+async def test_conversation_memory_initialization():
+    """Test that ConversationMemory initializes correctly."""
+    memory = ConversationMemory()
+
+    assert await memory.size() == 0
+    assert memory.messages == []
+    print("✅ ConversationMemory initialization works")
+
+
+@pytest.mark.asyncio
+async def test_add_single_message():
+    """Test adding a single message to memory."""
+    memory = ConversationMemory()
+
+    message = {"role": "user", "content": "Hello, world!"}
+    await memory.add(message)
+
+    assert await memory.size() == 1
+    messages = await memory.get_messages()
+    assert len(messages) == 1
+    assert messages[0]["role"] == "user"
+    assert messages[0]["content"] == "Hello, world!"
+    assert "id" in messages[0]  # Auto-generated ID
+    print("✅ Adding single message works")
+
+
+@pytest.mark.asyncio
+async def test_add_multiple_messages():
+    """Test adding multiple messages at once."""
+    memory = ConversationMemory()
+
+    messages = [
+        {"role": "user", "content": "Hello"},
+        {"role": "assistant", "content": "Hi there!"},
+        {"role": "user", "content": "How are you?"},
+    ]
+    await memory.add(messages)
+
+    assert await memory.size() == 3
+    retrieved = await memory.get_messages()
+    assert len(retrieved) == 3
+    assert retrieved[0]["content"] == "Hello"
+    assert retrieved[1]["content"] == "Hi there!"
+    assert retrieved[2]["content"] == "How are you?"
+    print("✅ Adding multiple messages works")
+
+
+@pytest.mark.asyncio
+async def test_duplicate_prevention():
+    """Test that duplicate messages are prevented by default."""
+    memory = ConversationMemory()
+
+    message = {"id": "msg_123", "role": "user", "content": "Hello"}
+
+    # Add same message twice
+    await memory.add(message)
+    await memory.add(message)
+
+    # Should only have one message
+    assert await memory.size() == 1
+    print("✅ Duplicate prevention works")
+
+
+@pytest.mark.asyncio
+async def test_allow_duplicates():
+    """Test allowing duplicate messages when explicitly enabled."""
+    memory = ConversationMemory()
+
+    message = {"id": "msg_123", "role": "user", "content": "Hello"}
+
+    # Add same message twice with allow_duplicates=True
+    await memory.add(message, allow_duplicates=True)
+    await memory.add(message, allow_duplicates=True)
+
+    # Should have two messages
+    assert await memory.size() == 2
+    print("✅ Allowing duplicates works")
+
+
+@pytest.mark.asyncio
+async def test_delete_by_index():
+    """Test deleting messages by index."""
+    memory = ConversationMemory()
+
+    messages = [
+        {"role": "user", "content": "Message 1"},
+        {"role": "assistant", "content": "Message 2"},
+        {"role": "user", "content": "Message 3"},
+    ]
+    await memory.add(messages)
+
+    # Delete middle message
+    await memory.delete(1)
+
+    assert await memory.size() == 2
+    retrieved = await memory.get_messages()
+    assert retrieved[0]["content"] == "Message 1"
+    assert retrieved[1]["content"] == "Message 3"
+    print("✅ Deleting by index works")
+
+
+@pytest.mark.asyncio
+async def test_delete_multiple_indices():
+    """Test deleting multiple messages at once."""
+    memory = ConversationMemory()
+
+    messages = [{"role": "user", "content": f"Message {i}"} for i in range(5)]
+    await memory.add(messages)
+
+    # Delete indices 1 and 3
+    await memory.delete([1, 3])
+
+    assert await memory.size() == 3
+    retrieved = await memory.get_messages()
+    assert retrieved[0]["content"] == "Message 0"
+    assert retrieved[1]["content"] == "Message 2"
+    assert retrieved[2]["content"] == "Message 4"
+    print("✅ Deleting multiple indices works")
+
+
+@pytest.mark.asyncio
+async def test_delete_invalid_index():
+    """Test that deleting invalid index raises error."""
+    memory = ConversationMemory()
+
+    await memory.add({"role": "user", "content": "Hello"})
+
+    # Try to delete out of range index
+    with pytest.raises(IndexError):
+        await memory.delete(10)
+
+    print("✅ Invalid index deletion raises error correctly")
+
+
+@pytest.mark.asyncio
+async def test_get_last_message():
+    """Test getting the last message."""
+    memory = ConversationMemory()
+
+    # Empty memory
+    assert await memory.get_last_message() is None
+
+    # Add messages
+    messages = [
+        {"role": "user", "content": "First"},
+        {"role": "assistant", "content": "Second"},
+        {"role": "user", "content": "Third"},
+    ]
+    await memory.add(messages)
+
+    last = await memory.get_last_message()
+    assert last is not None
+    assert last["content"] == "Third"
+    print("✅ Getting last message works")
+
+
+@pytest.mark.asyncio
+async def test_get_messages_by_role():
+    """Test filtering messages by role."""
+    memory = ConversationMemory()
+
+    messages = [
+        {"role": "user", "content": "User 1"},
+        {"role": "assistant", "content": "Assistant 1"},
+        {"role": "user", "content": "User 2"},
+        {"role": "assistant", "content": "Assistant 2"},
+        {"role": "system", "content": "System 1"},
+    ]
+    await memory.add(messages)
+
+    user_messages = await memory.get_messages_by_role("user")
+    assert len(user_messages) == 2
+    assert all(msg["role"] == "user" for msg in user_messages)
+
+    assistant_messages = await memory.get_messages_by_role("assistant")
+    assert len(assistant_messages) == 2
+    assert all(msg["role"] == "assistant" for msg in assistant_messages)
+
+    system_messages = await memory.get_messages_by_role("system")
+    assert len(system_messages) == 1
+
+    print("✅ Filtering by role works")
+
+
+@pytest.mark.asyncio
+async def test_get_messages_with_limit():
+    """Test getting messages with limit."""
+    memory = ConversationMemory()
+
+    messages = [{"role": "user", "content": f"Message {i}"} for i in range(10)]
+    await memory.add(messages)
+
+    # Get last 3 messages
+    recent = await memory.get_messages(limit=3)
+    assert len(recent) == 3
+    assert recent[0]["content"] == "Message 7"
+    assert recent[1]["content"] == "Message 8"
+    assert recent[2]["content"] == "Message 9"
+    print("✅ Getting messages with limit works")
+
+
+@pytest.mark.asyncio
+async def test_truncate_to_size():
+    """Test truncating memory to a maximum size."""
+    memory = ConversationMemory()
+
+    messages = [{"role": "user", "content": f"Message {i}"} for i in range(10)]
+    await memory.add(messages)
+
+    # Truncate to last 5 messages
+    await memory.truncate_to_size(5)
+
+    assert await memory.size() == 5
+    retrieved = await memory.get_messages()
+    assert retrieved[0]["content"] == "Message 5"
+    assert retrieved[4]["content"] == "Message 9"
+    print("✅ Truncating to size works")
+
+
+@pytest.mark.asyncio
+async def test_clear_memory():
+    """Test clearing all messages from memory."""
+    memory = ConversationMemory()
+
+    messages = [{"role": "user", "content": f"Message {i}"} for i in range(5)]
+    await memory.add(messages)
+
+    assert await memory.size() == 5
+
+    await memory.clear()
+
+    assert await memory.size() == 0
+    assert await memory.get_messages() == []
+    print("✅ Clearing memory works")
+
+
+@pytest.mark.asyncio
+async def test_state_dict_serialization():
+    """Test state serialization and deserialization."""
+    memory1 = ConversationMemory()
+
+    messages = [
+        {"role": "user", "content": "Hello"},
+        {"role": "assistant", "content": "Hi!"},
+    ]
+    await memory1.add(messages)
+
+    # Export state
+    state = memory1.state_dict()
+    assert "messages" in state
+    assert len(state["messages"]) == 2
+
+    # Load into new memory
+    memory2 = ConversationMemory()
+    memory2.load_state_dict(state)
+
+    assert await memory2.size() == 2
+    retrieved = await memory2.get_messages()
+    assert retrieved[0]["content"] == "Hello"
+    assert retrieved[1]["content"] == "Hi!"
+    print("✅ State serialization works")
+
+
+@pytest.mark.asyncio
+async def test_state_dict_strict_mode():
+    """Test state loading with strict mode validation."""
+    memory = ConversationMemory()
+
+    # Invalid state dict (missing 'messages' key)
+    invalid_state = {"wrong_key": []}
+
+    # Should raise error in strict mode
+    with pytest.raises(ValueError):
+        memory.load_state_dict(invalid_state, strict=True)
+
+    # Should not raise error in non-strict mode
+    memory.load_state_dict(invalid_state, strict=False)
+    assert await memory.size() == 0
+
+    print("✅ State dict strict mode works")
+
+
+@pytest.mark.asyncio
+async def test_add_none_message():
+    """Test that adding None message is handled gracefully."""
+    memory = ConversationMemory()
+
+    await memory.add(None)
+    assert await memory.size() == 0
+    print("✅ Adding None message handled gracefully")
+
+
+@pytest.mark.asyncio
+async def test_add_invalid_message_type():
+    """Test that adding invalid message type raises error."""
+    memory = ConversationMemory()
+
+    # Try to add a string instead of dict
+    with pytest.raises(TypeError):
+        await memory.add("invalid message")
+
+    # Try to add list of non-dicts
+    with pytest.raises(TypeError):
+        await memory.add(["invalid", "messages"])
+
+    print("✅ Invalid message type raises error correctly")
+
+
+@pytest.mark.asyncio
+async def test_retrieve_not_implemented():
+    """Test that retrieve method raises NotImplementedError."""
+    memory = ConversationMemory()
+
+    with pytest.raises(NotImplementedError):
+        await memory.retrieve("some query")
+
+    print("✅ Retrieve method correctly not implemented")
+
+
+@pytest.mark.asyncio
+async def test_message_isolation():
+    """Test that returned messages are copies, not references."""
+    memory = ConversationMemory()
+
+    original = {"role": "user", "content": "Original"}
+    await memory.add(original)
+
+    # Get messages and modify
+    retrieved = await memory.get_messages()
+    retrieved[0]["content"] = "Modified"
+
+    # Original in memory should be unchanged
+    messages = await memory.get_messages()
+    assert messages[0]["content"] == "Original"
+    print("✅ Message isolation works correctly")
+
+
+if __name__ == "__main__":
+    import asyncio
+
+    async def run_all_tests():
+        """Run all tests manually."""
+        print("\n=== Running ConversationMemory Tests ===\n")
+
+        await test_conversation_memory_initialization()
+        await test_add_single_message()
+        await test_add_multiple_messages()
+        await test_duplicate_prevention()
+        await test_allow_duplicates()
+        await test_delete_by_index()
+        await test_delete_multiple_indices()
+        await test_delete_invalid_index()
+        await test_get_last_message()
+        await test_get_messages_by_role()
+        await test_get_messages_with_limit()
+        await test_truncate_to_size()
+        await test_clear_memory()
+        await test_state_dict_serialization()
+        await test_state_dict_strict_mode()
+        await test_add_none_message()
+        await test_add_invalid_message_type()
+        await test_retrieve_not_implemented()
+        await test_message_isolation()
+
+        print("\n=== All ConversationMemory Tests Passed! ===\n")
+
+    asyncio.run(run_all_tests())

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"])