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

massgen/tests/test_rich_terminal_display.py

@@ -1,4 +1,5 @@
 #!/usr/bin/env python3
+# -*- coding: utf-8 -*-
 """
 Test script for MassGen Rich Terminal Display.
 Tests RichTerminalDisplay functionality with two-agent coordination.
@@ -13,11 +14,16 @@ from pathlib import Path
 project_root = Path(__file__).parent.parent.parent
 sys.path.insert(0, str(project_root))
 
-from massgen.backend.response import ResponseBackend
-from massgen.chat_agent import SingleAgent
-from massgen.orchestrator import Orchestrator
-from massgen.frontend.coordination_ui import CoordinationUI, coordinate_with_rich_ui
-from massgen.frontend.displays.rich_terminal_display import is_rich_available
+from massgen.backend.response import ResponseBackend  # noqa: E402
+from massgen.chat_agent import SingleAgent  # noqa: E402
+from massgen.frontend.coordination_ui import (  # noqa: E402
+    CoordinationUI,
+    coordinate_with_rich_ui,
+)
+from massgen.frontend.displays.rich_terminal_display import (  # noqa: E402
+    is_rich_available,
+)
+from massgen.orchestrator import Orchestrator  # noqa: E402
 
 
 async def test_rich_availability():
@@ -64,7 +70,7 @@ async def test_rich_display_basic():
         # Test theme configuration
         themes = ["dark", "light", "cyberpunk"]
         for theme in themes:
-            themed_display = RichTerminalDisplay(agent_ids, theme=theme)
+            RichTerminalDisplay(agent_ids, theme=theme)
             print(f"✅ {theme.title()} theme created successfully")
 
         # Clean up
@@ -186,9 +192,7 @@ async def test_rich_convenience_function():
             system_message="You are a strategic thinker who focuses on long-term implications and strategic recommendations.",
         )
 
-        orchestrator = Orchestrator(
-            agents={"analyst": analyst, "strategist": strategist}
-        )
+        orchestrator = Orchestrator(agents={"analyst": analyst, "strategist": strategist})
 
         print("🎯 Testing convenience function with light theme...")
 
@@ -231,13 +235,10 @@ async def test_rich_fallback():
 
     try:
         # Test UI creation with rich_terminal when Rich is available
-        ui = CoordinationUI(display_type="rich_terminal")
+        CoordinationUI(display_type="rich_terminal")
 
         if is_rich_available():
             print("✅ Rich is available - RichTerminalDisplay should be used")
-            from massgen.frontend.displays.rich_terminal_display import (
-                RichTerminalDisplay,
-            )
 
             # Note: We can't easily test the actual fallback without mocking
             print("📝 Note: Fallback logic tested through UI creation")

massgen/tests/test_timeout.py

@@ -0,0 +1,133 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Simple test script to verify timeout mechanisms work correctly.
+This creates a fast timeout config and runs a test to demonstrate the timeout fallback.
+"""
+
+import asyncio
+import sys
+from pathlib import Path
+
+from massgen.agent_config import AgentConfig, TimeoutConfig
+from massgen.backend.claude_code import ClaudeCodeBackend
+from massgen.chat_agent import SingleAgent
+from massgen.orchestrator import Orchestrator
+
+# Add massgen to Python path
+sys.path.insert(0, str(Path(__file__).parent))
+
+
+async def test_orchestrator_timeout():
+    """Test orchestrator-level timeout."""
+    print("\n🧪 Testing Orchestrator Timeout Mechanism")
+    print("=" * 50)
+
+    # Create very restrictive orchestrator timeout
+    timeout_config = TimeoutConfig(
+        orchestrator_timeout_seconds=10,  # 10 seconds orchestrator (very short)
+    )
+
+    # Create agent config and set timeout
+    agent_config = AgentConfig.create_claude_code_config(model="claude-sonnet-4-20250514")
+    agent_config.timeout_config = timeout_config
+
+    # Claude code backend for testing
+    try:
+        backend = ClaudeCodeBackend()
+        agent = SingleAgent(backend=backend, system_message="You are a helpful assistant.")
+
+        # Create orchestrator with timeout-aware agents
+        agents = {"test_agent": agent}
+        orchestrator = Orchestrator(agents=agents, config=agent_config)
+
+        print(f"⏱️  Orchestrator timeout: {timeout_config.orchestrator_timeout_seconds}s")
+        print("📝 Testing with complex multi-agent coordination that should trigger orchestrator timeout...")
+
+        # Ask a question that requires complex coordination between multiple agents
+        question = (
+            "Please coordinate with multiple specialized agents to create a comprehensive business plan "
+            "for a tech startup, including market analysis, technical architecture, financial projections, "
+            "legal considerations, and detailed implementation timeline. Each section should be thoroughly "
+            "researched and cross-validated between agents."
+        )
+
+        print(f"\n❓ Question: {question[:100]}...")
+        print("\n🚀 Starting orchestrator coordination (should timeout quickly)...")
+
+        response_content = ""
+        timeout_detected = False
+
+        async for chunk in orchestrator.chat_simple(question):
+            if chunk.type == "content":
+                content = chunk.content
+                print(f"📝 {content}")
+                response_content += chunk.content
+                if "time limit exceeded" in content.lower() or "timeout" in content.lower():
+                    timeout_detected = True
+                    print(f"⚠️  ORCHESTRATOR TIMEOUT DETECTED: {content}")
+            elif chunk.type == "error":
+                if "time limit exceeded" in chunk.error.lower() or "timeout" in chunk.error.lower():
+                    timeout_detected = True
+                    print(f"⚠️  ORCHESTRATOR TIMEOUT DETECTED: {chunk.error}")
+            elif chunk.type == "done":
+                print("✅ Orchestrator coordination completed")
+                break
+
+        if timeout_detected:
+            print("\n🎯 SUCCESS: Orchestrator timeout mechanism triggered correctly!")
+        else:
+            print("\n🤔 No orchestrator timeout detected - either coordination completed fast or timeout didn't work")
+
+        print(f"\n📊 Final response length: {len(response_content)} characters")
+
+    except Exception as e:
+        print(f"❌ Orchestrator timeout test failed with error: {e}")
+        print("💡 Note: This test requires API keys to run with real backend")
+
+
+def print_config_example():
+    """Print example configuration for users."""
+    print("\n📋 Example YAML Configuration with Timeout Settings:")
+    print("=" * 50)
+
+    example_config = """
+# Conservative timeout settings to prevent runaway costs
+timeout_settings:
+  orchestrator_timeout_seconds: 600   # 10 minutes max coordination
+
+agents:
+  - id: "agent1"
+    backend:
+      type: "openai"
+      model: "gpt-4o-mini"
+    system_message: "You are a helpful assistant."
+"""
+
+    print(example_config)
+
+    print("\n🖥️  CLI Examples:")
+    print('python -m massgen.cli --config config.yaml --orchestrator-timeout 300 "Complex task"')
+
+
+if __name__ == "__main__":
+    print("🔧 MassGen Timeout Mechanism Test")
+    print("=" * 60)
+
+    print_config_example()
+
+    print("\n🧪 Running timeout tests...")
+    print("Note: These tests require API keys to run with real backends")
+
+    try:
+        # Run orchestrator timeout test
+        asyncio.run(test_orchestrator_timeout())
+
+        print("\n✅ Timeout mechanism implementation completed!")
+        print("💡 The timeout system will prevent runaway token usage and provide graceful fallbacks.")
+
+    except KeyboardInterrupt:
+        print("\n⏹️  Tests interrupted by user")
+    except Exception as e:
+        print(f"\n❌ Tests failed: {e}")
+        print("💡 This is expected if you don't have API keys configured")

massgen/tests/test_v3_3agents.py

@@ -1,9 +1,10 @@
 #!/usr/bin/env python3
+# -*- coding: utf-8 -*-
 """
 MassGen Example: Three Agent Coordination
 
-This example demonstrates three-agent coordination using the 
-multi-region coordination UI. Three agents with different specialties work 
+This example demonstrates three-agent coordination using the
+multi-region coordination UI. Three agents with different specialties work
 together on a question that benefits from multiple perspectives.
 
 Features:
@@ -16,20 +17,20 @@ Note: Requires OPENAI_API_KEY in environment.
 """
 
 import asyncio
-import sys
 import os
+import sys
+
+from massgen import Orchestrator, ResponseBackend, create_simple_agent
+from massgen.frontend.coordination_ui import coordinate_with_terminal_ui
 
 # Add project root to path
 project_root = os.path.dirname(__file__)
 sys.path.insert(0, project_root)
 
-from massgen import create_simple_agent, ResponseBackend, Orchestrator
-from massgen.frontend.coordination_ui import coordinate_with_terminal_ui
-
 
 async def three_agent_example():
     """Demonstrate three agent coordination with multi-region UI."""
-    
+
     print("🎯 MassGen: Three Agent Coordination")
     print("=" * 50)
 
@@ -65,9 +66,7 @@ async def three_agent_example():
         print("  • educator: Clear explanations")
 
         # Create orchestrator
-        orchestrator = Orchestrator(
-            agents={"scientist": scientist, "engineer": engineer, "educator": educator}
-        )
+        orchestrator = Orchestrator(agents={"scientist": scientist, "engineer": engineer, "educator": educator})
         print("✅ orchestrator ready")
 
         # Question that benefits from multiple perspectives
@@ -83,14 +82,14 @@ async def three_agent_example():
             logging_enabled=False,
         )
 
-        print(f"\n✅ Three agent coordination completed!")
+        print("\n✅ Three agent coordination completed!")
 
         # Show results
         orchestrator_status = orchestrator.get_status()
         selected_agent = orchestrator_status.get("selected_agent")
         vote_results = orchestrator_status.get("vote_results", {})
 
-        print(f"\n📊 Results:")
+        print("\n📊 Results:")
         print(f"🏆 Selected agent: {selected_agent}")
         if vote_results.get("vote_counts"):
             for agent, votes in vote_results["vote_counts"].items():

massgen/tests/test_v3_simple.py

@@ -1,4 +1,5 @@
 #!/usr/bin/env python3
+# -*- coding: utf-8 -*-
 """
 Simple test script for MassGen basic functionality.
 Tests single agent and basic orchestrator functionality.
@@ -9,14 +10,14 @@ import os
 import sys
 from pathlib import Path
 
-# Add project root to path
-project_root = Path(__file__).parent.parent.parent.parent
-sys.path.insert(0, str(project_root))
-
 from massgen.backend.response import ResponseBackend
 from massgen.chat_agent import SingleAgent
 from massgen.orchestrator import Orchestrator
 
+# Add project root to path
+project_root = Path(__file__).parent.parent.parent.parent
+sys.path.insert(0, str(project_root))
+
 
 async def test_single_agent():
     """Test basic SingleAgent functionality."""
@@ -51,9 +52,7 @@ async def test_single_agent():
             elif chunk.type == "done":
                 break
 
-        print(
-            f"\n✅ Single agent test completed. Response: '{response_content.strip()}'"
-        )
+        print(f"\n✅ Single agent test completed. Response: '{response_content.strip()}'")
         return True
 
     except Exception as e:
@@ -83,9 +82,7 @@ async def test_orchestrator_single():
 
         print("📤 Testing orchestrator with single agent...")
 
-        messages = [
-            {"role": "user", "content": "Explain photosynthesis in one sentence."}
-        ]
+        messages = [{"role": "user", "content": "Explain photosynthesis in one sentence."}]
         response_content = ""
 
         async for chunk in orchestrator.chat(messages):
@@ -98,9 +95,7 @@ async def test_orchestrator_single():
             elif chunk.type == "done":
                 break
 
-        print(
-            f"\n✅ Orchestrator single agent test completed. Response length: {len(response_content)} chars"
-        )
+        print(f"\n✅ Orchestrator single agent test completed. Response length: {len(response_content)} chars")
         return True
 
     except Exception as e:

massgen/tests/test_v3_three_agents.py

@@ -1,4 +1,5 @@
 #!/usr/bin/env python3
+# -*- coding: utf-8 -*-
 """
 Test script for MassGen three-agent coordination with terminal display.
 Tests orchestrator coordination between three agents with diverse expertise.
@@ -13,10 +14,10 @@ from pathlib import Path
 project_root = Path(__file__).parent.parent.parent.parent
 sys.path.insert(0, str(project_root))
 
-from massgen.backend.response import ResponseBackend
-from massgen.chat_agent import SingleAgent
-from massgen.orchestrator import Orchestrator
-from massgen.frontend.coordination_ui import CoordinationUI
+from massgen.backend.response import ResponseBackend  # noqa: E402
+from massgen.chat_agent import SingleAgent  # noqa: E402
+from massgen.frontend.coordination_ui import CoordinationUI  # noqa: E402
+from massgen.orchestrator import Orchestrator  # noqa: E402
 
 
 async def test_three_agents_coordination():
@@ -128,9 +129,7 @@ async def test_three_agents_simple():
             system_message="You are an evaluator. Focus on assessment and quality control.",
         )
 
-        orchestrator = Orchestrator(
-            agents={"planner": planner, "executor": executor, "evaluator": evaluator}
-        )
+        orchestrator = Orchestrator(agents={"planner": planner, "executor": executor, "evaluator": evaluator})
 
         print("📤 Testing simple three-agent coordination...")
 
@@ -138,7 +137,7 @@ async def test_three_agents_simple():
             {
                 "role": "user",
                 "content": "How can we improve team productivity in a remote work environment?",
-            }
+            },
         ]
 
         response_content = ""
@@ -152,9 +151,7 @@ async def test_three_agents_simple():
             elif chunk.type == "done":
                 break
 
-        print(
-            f"\n✅ Simple test completed. Response length: {len(response_content)} characters"
-        )
+        print(f"\n✅ Simple test completed. Response length: {len(response_content)} characters")
         return True
 
     except Exception as e:
@@ -194,9 +191,7 @@ async def test_three_agents_consensus():
             system_message="You are a skeptical viewpoint agent. Focus on potential challenges and risks.",
         )
 
-        orchestrator = Orchestrator(
-            agents={"optimist": optimist, "realist": realist, "skeptic": skeptic}
-        )
+        orchestrator = Orchestrator(agents={"optimist": optimist, "realist": realist, "skeptic": skeptic})
 
         print("📤 Testing consensus building with diverse viewpoints...")
 
@@ -205,7 +200,7 @@ async def test_three_agents_consensus():
             {
                 "role": "user",
                 "content": "What is the future outlook for artificial intelligence in healthcare?",
-            }
+            },
         ]
 
         response_content = ""
@@ -219,9 +214,7 @@ async def test_three_agents_consensus():
             elif chunk.type == "done":
                 break
 
-        print(
-            f"\n✅ Consensus test completed. Response length: {len(response_content)} characters"
-        )
+        print(f"\n✅ Consensus test completed. Response length: {len(response_content)} characters")
         return True
 
     except Exception as e:

massgen/tests/test_v3_two_agents.py

@@ -1,4 +1,5 @@
 #!/usr/bin/env python3
+# -*- coding: utf-8 -*-
 """
 Test script for MassGen two-agent coordination with terminal display.
 Tests orchestrator coordination between two agents with different expertise.
@@ -13,10 +14,10 @@ from pathlib import Path
 project_root = Path(__file__).parent.parent.parent.parent
 sys.path.insert(0, str(project_root))
 
-from massgen.backend.response import ResponseBackend
-from massgen.chat_agent import SingleAgent
-from massgen.orchestrator import Orchestrator
-from massgen.frontend.coordination_ui import CoordinationUI
+from massgen.backend.response import ResponseBackend  # noqa: E402
+from massgen.chat_agent import SingleAgent  # noqa: E402
+from massgen.frontend.coordination_ui import CoordinationUI  # noqa: E402
+from massgen.orchestrator import Orchestrator  # noqa: E402
 
 
 async def test_two_agents_coordination():
@@ -62,9 +63,7 @@ async def test_two_agents_coordination():
         print()
 
         # Test question that benefits from both perspectives
-        test_question = (
-            "How does photosynthesis work and why is it important for life on Earth?"
-        )
+        test_question = "How does photosynthesis work and why is it important for life on Earth?"
 
         print(f"📝 Question: {test_question}")
         print("\n🎭 Starting two-agent coordination...")
@@ -117,9 +116,7 @@ async def test_two_agents_simple():
 
         print("📤 Testing simple coordination...")
 
-        messages = [
-            {"role": "user", "content": "What are the benefits of renewable energy?"}
-        ]
+        messages = [{"role": "user", "content": "What are the benefits of renewable energy?"}]
 
         response_content = ""
         async for chunk in orchestrator.chat(messages):
@@ -132,9 +129,7 @@ async def test_two_agents_simple():
             elif chunk.type == "done":
                 break
 
-        print(
-            f"\n✅ Simple test completed. Response length: {len(response_content)} characters"
-        )
+        print(f"\n✅ Simple test completed. Response length: {len(response_content)} characters")
         return True
 
     except Exception as e:

massgen/token_manager/__init__.py

@@ -0,0 +1,7 @@
+# -*- coding: utf-8 -*-
+from .token_manager import TokenCostCalculator, TokenUsage
+
+__all__ = [
+    "TokenUsage",
+    "TokenCostCalculator",
+]