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

massgen/tests/test_chat_completions_refactor.py

@@ -0,0 +1,142 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Test script to verify ChatCompletionsBackend refactoring.
+Tests integration with different OpenAI-compatible providers.
+"""
+
+import asyncio
+import os
+
+from massgen.backend import ChatCompletionsBackend
+
+
+async def test_openai_backend():
+    """Test ChatCompletionsBackend with OpenAI."""
+    print("🔧 Testing ChatCompletionsBackend with OpenAI...")
+
+    # Create backend with OpenAI defaults
+    backend = ChatCompletionsBackend()
+    print(f"Provider: {backend.get_provider_name()}")
+    print(f"Base URL: {backend.base_url}")
+    print(f"API Key configured: {'Yes' if backend.api_key else 'No'}")
+
+    # Test token estimation
+    test_text = "Hello world, how are you doing today?"
+    tokens = backend.estimate_tokens(test_text)
+    print(f"Estimated tokens for '{test_text}': {tokens}")
+
+    # Test cost calculation
+    cost = backend.calculate_cost(1000, 500, "gpt-4o-mini")
+    print(f"Cost for 1000 input + 500 output tokens (gpt-4o-mini): ${cost:.4f}")
+
+
+async def test_together_ai_backend():
+    """Test ChatCompletionsBackend configured for Together AI."""
+    print("\n🔧 Testing ChatCompletionsBackend with Together AI...")
+
+    # Create backend configured for Together AI
+    backend = ChatCompletionsBackend(
+        base_url="https://api.together.xyz/v1",
+        provider_name="Together AI",
+        api_key=os.getenv("TOGETHER_API_KEY"),
+    )
+    print(f"Provider: {backend.get_provider_name()}")
+    print(f"Base URL: {backend.base_url}")
+    print(f"API Key configured: {'Yes' if backend.api_key else 'No'}")
+
+    # Test cost calculation with Together AI model
+    cost = backend.calculate_cost(1000, 500, "meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo")
+    print(f"Cost for 1000 input + 500 output tokens (fallback pricing): ${cost:.4f}")
+
+
+async def test_cerebras_backend():
+    """Test ChatCompletionsBackend configured for Cerebras AI."""
+    print("\n🔧 Testing ChatCompletionsBackend with Cerebras AI...")
+
+    # Create backend configured for Cerebras AI
+    backend = ChatCompletionsBackend(
+        base_url="https://api.cerebras.ai/v1",
+        provider_name="Cerebras AI",
+        api_key=os.getenv("CEREBRAS_API_KEY"),
+    )
+    print(f"Provider: {backend.get_provider_name()}")
+    print(f"Base URL: {backend.base_url}")
+    print(f"API Key configured: {'Yes' if backend.api_key else 'No'}")
+
+
+async def test_tool_conversion():
+    """Test tool format conversion."""
+    print("\n🔧 Testing tool format conversion...")
+
+    backend = ChatCompletionsBackend()
+
+    # Test Response API format conversion
+    response_tools = [
+        {
+            "type": "function",
+            "name": "get_weather",
+            "description": "Get weather information",
+            "parameters": {
+                "type": "object",
+                "properties": {"location": {"type": "string"}},
+                "required": ["location"],
+            },
+        },
+    ]
+
+    converted = backend.convert_tools_to_chat_completions_format(response_tools)
+    print("Response API tools converted to Chat Completions format:")
+    print(f"  Original: {response_tools[0]}")
+    print(f"  Converted: {converted[0]}")
+
+    # Test Chat Completions format (should remain unchanged)
+    chat_tools = [
+        {
+            "type": "function",
+            "function": {
+                "name": "search_web",
+                "description": "Search the web",
+                "parameters": {
+                    "type": "object",
+                    "properties": {"query": {"type": "string"}},
+                    "required": ["query"],
+                },
+            },
+        },
+    ]
+
+    converted_chat = backend.convert_tools_to_chat_completions_format(chat_tools)
+    print("\nChat Completions tools (should remain unchanged):")
+    print(f"  Original: {chat_tools[0]}")
+    print(f"  After conversion: {converted_chat[0]}")
+
+
+async def main():
+    """Run all tests."""
+    print("🚀 Testing ChatCompletionsBackend refactoring...")
+    print("=" * 60)
+
+    try:
+        await test_openai_backend()
+        await test_together_ai_backend()
+        await test_cerebras_backend()
+        await test_tool_conversion()
+
+        print("\n" + "=" * 60)
+        print("✅ All tests completed successfully!")
+        print("\n📋 ChatCompletionsBackend is now ready for use with:")
+        print("   • OpenAI (default)")
+        print("   • Together AI")
+        print("   • Cerebras AI")
+        print("   • Any OpenAI-compatible provider")
+
+    except Exception as e:
+        print(f"\n❌ Test failed: {e}")
+        return 1
+
+    return 0
+
+
+if __name__ == "__main__":
+    exit(asyncio.run(main()))

massgen/tests/test_claude_backend.py

@@ -1,4 +1,5 @@
 #!/usr/bin/env python3
+# -*- coding: utf-8 -*-
 """
 Claude Backend Integration Tests for MassGen
 
@@ -12,18 +13,16 @@ Tests the Claude backend implementation with real API calls:
 Requires ANTHROPIC_API_KEY environment variable.
 """
 
+import asyncio
 import os
 import sys
-import asyncio
-import json
 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.claude_backend import ClaudeBackend
-from massgen.backend.base import StreamChunk
+from massgen.backend.claude import ClaudeBackend  # noqa: E402
 
 
 async def test_claude_basic_streaming():
@@ -32,21 +31,15 @@ async def test_claude_basic_streaming():
 
     backend = ClaudeBackend()
 
-    messages = [
-        {"role": "user", "content": "Explain quantum computing in 2-3 sentences."}
-    ]
+    messages = [{"role": "user", "content": "Explain quantum computing in 2-3 sentences."}]
 
     content = ""
-    async for chunk in backend.stream_with_tools(
-        messages, [], model="claude-3-5-haiku-20241022"
-    ):
+    async for chunk in backend.stream_with_tools(messages, [], model="claude-3-5-haiku-20241022"):
         if chunk.type == "content":
             content += chunk.content
             print(chunk.content, end="", flush=True)
         elif chunk.type == "complete_message":
-            print(
-                f"\n✅ Complete message received: {len(chunk.complete_message.get('content', ''))} chars"
-            )
+            print(f"\n✅ Complete message received: {len(chunk.complete_message.get('content', ''))} chars")
         elif chunk.type == "done":
             print("\n✅ Basic streaming test completed")
             break
@@ -77,20 +70,18 @@ async def test_claude_tool_calling():
                 },
                 "required": ["width", "height"],
             },
-        }
+        },
     ]
 
     messages = [
         {
             "role": "user",
             "content": "Calculate the area of a rectangle with width 5 and height 3.",
-        }
+        },
     ]
 
     tool_calls_received = []
-    async for chunk in backend.stream_with_tools(
-        messages, tools, model="claude-3-5-haiku-20241022"
-    ):
+    async for chunk in backend.stream_with_tools(messages, tools, model="claude-3-5-haiku-20241022"):
         if chunk.type == "content":
             print(chunk.content, end="", flush=True)
         elif chunk.type == "tool_calls":
@@ -101,7 +92,7 @@ async def test_claude_tool_calling():
                 tool_args = backend.extract_tool_arguments(tool_call)
                 print(f"   - {tool_name}: {tool_args}")
         elif chunk.type == "complete_message":
-            print(f"\n✅ Complete message with tool calls received")
+            print("\n✅ Complete message with tool calls received")
         elif chunk.type == "done":
             print("✅ Tool calling test completed")
             break
@@ -129,14 +120,14 @@ async def test_claude_multi_tool_support():
                 "properties": {"title": {"type": "string"}, "data": {"type": "string"}},
                 "required": ["title", "data"],
             },
-        }
+        },
     ]
 
     messages = [
         {
             "role": "user",
             "content": "Search for recent news about AI and format the result with a nice title.",
-        }
+        },
     ]
 
     # Enable both server-side tools and user tools
@@ -187,7 +178,7 @@ async def test_claude_message_conversion():
                     "id": "call_123",
                     "type": "function",
                     "function": {"name": "add", "arguments": {"a": 5, "b": 3}},
-                }
+                },
             ],
         },
         {"role": "tool", "tool_call_id": "call_123", "content": "8"},
@@ -203,11 +194,7 @@ async def test_claude_message_conversion():
     # Check tool result conversion
     tool_result_found = False
     for msg in converted:
-        if (
-            msg.get("role") == "user"
-            and isinstance(msg.get("content"), list)
-            and any(item.get("type") == "tool_result" for item in msg["content"])
-        ):
+        if msg.get("role") == "user" and isinstance(msg.get("content"), list) and any(item.get("type") == "tool_result" for item in msg["content"]):
             tool_result_found = True
             print("   ✅ Tool result conversion successful")
             break
@@ -290,7 +277,7 @@ async def main():
             results.append((test_name, False))
 
     # Summary
-    print(f"\n📊 Test Results Summary:")
+    print("\n📊 Test Results Summary:")
     passed = sum(1 for _, result in results if result)
     total = len(results)
 

massgen/tests/test_claude_code.py

@@ -0,0 +1,268 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Real test of ClaudeCodeBackend with actual Claude Code API calls.
+This test outputs the actual stream chunks to verify functionality.
+"""
+
+import asyncio
+import os
+
+from massgen.backend.claude_code import ClaudeCodeBackend
+
+
+async def test_real_stream_with_tools():
+    """Test real streaming with Claude Code API and output stream chunks."""
+
+    # Check if API key is available
+    if not os.getenv("ANTHROPIC_API_KEY"):
+        print("❌ ANTHROPIC_API_KEY not found in environment")
+        print("Set your API key: export ANTHROPIC_API_KEY=your-key-here")
+        return
+
+    print("🚀 Testing ClaudeCodeBackend with real Claude Code API")
+    print("=" * 60)
+
+    # Initialize backend
+    backend = ClaudeCodeBackend()
+    print(f"✅ Backend initialized: {backend.get_provider_name()}")
+    print(f"📊 Supported tools: {len(backend.get_supported_builtin_tools())} tools")
+
+    # Test single turn conversation
+    print("\n🔄 Testing single turn conversation...")
+    messages = [
+        {
+            "role": "user",
+            "content": "Hello! Can you tell me what 2+2 equals and show your calculation?",
+        },
+    ]
+
+    try:
+        print("\n📨 Sending messages:", messages)
+        print("\n📡 Stream chunks received:")
+        print("-" * 40)
+
+        chunk_count = 0
+        total_content = ""
+
+        async for chunk in backend.stream_with_tools(messages, []):
+            chunk_count += 1
+            print(f"[{chunk_count:2d}] Type: {chunk.type:<20} Source: {chunk.source or 'None':<20}")
+
+            if chunk.type == "content":
+                print(f"     Content: {repr(chunk.content)}")
+                total_content += chunk.content or ""
+            elif chunk.type == "complete_message":
+                print(f"     Complete message: {chunk.complete_message}")
+            elif chunk.type == "complete_response":
+                print(f"     Response metadata: {chunk.response}")
+            elif chunk.type == "agent_status":
+                print(f"     Status: {chunk.status} - {chunk.content}")
+            elif chunk.type == "builtin_tool_results":
+                print(f"     Tool results: {chunk.builtin_tool_results}")
+            elif chunk.type == "tool_calls":
+                print(f"     Tool calls: {chunk.tool_calls}")
+            elif chunk.type == "error":
+                print(f"     Error: {chunk.error}")
+            elif chunk.type == "done":
+                print("     ✅ Stream completed")
+                break
+
+            print()
+
+        print("-" * 40)
+        print(f"📊 Total chunks: {chunk_count}")
+        print(f"📝 Total content length: {len(total_content)} chars")
+        print(f"💰 Token usage: {backend.get_token_usage()}")
+        print(f"🔗 Session ID: {backend.get_current_session_id()}")
+
+        if total_content:
+            print(f"\n📄 Complete response:\n{total_content}")
+
+    except Exception as e:
+        print(f"❌ Error during streaming: {e}")
+        import traceback
+
+        traceback.print_exc()
+        return
+
+    # Test multi-turn conversation
+    print("\n" + "=" * 60)
+    print("🔄 Testing multi-turn conversation...")
+
+    # Add assistant response to conversation
+    messages.append({"role": "assistant", "content": total_content})
+
+    # Add follow-up question
+    messages.append(
+        {
+            "role": "user",
+            "content": "Great! Now can you show me how to calculate the result times 5?",
+        },
+    )
+
+    try:
+        print("\n📨 Sending multi-turn messages:")
+        for i, msg in enumerate(messages):
+            print(f"  [{i+1}] {msg['role']}: {msg['content'][:100]}{'...' if len(msg['content']) > 100 else ''}")
+
+        print("\n📡 Stream chunks received:")
+        print("-" * 40)
+
+        chunk_count = 0
+        turn2_content = ""
+
+        async for chunk in backend.stream_with_tools(messages, []):
+            chunk_count += 1
+            print(f"[{chunk_count:2d}] Type: {chunk.type:<20} Source: {chunk.source or 'None':<20}")
+
+            if chunk.type == "content":
+                print(f"     Content: {repr(chunk.content)}")
+                turn2_content += chunk.content or ""
+            elif chunk.type == "complete_response":
+                print(f"     Response metadata: {chunk.response}")
+            elif chunk.type == "done":
+                print("     ✅ Stream completed")
+                break
+
+            print()
+
+        print("-" * 40)
+        print(f"📊 Turn 2 chunks: {chunk_count}")
+        print(f"📝 Turn 2 content length: {len(turn2_content)} chars")
+        print(f"💰 Cumulative token usage: {backend.get_token_usage()}")
+        print(f"🔗 Session ID: {backend.get_current_session_id()}")
+
+        if turn2_content:
+            print(f"\n📄 Turn 2 response:\n{turn2_content}")
+
+    except Exception as e:
+        print(f"❌ Error during multi-turn: {e}")
+        import traceback
+
+        traceback.print_exc()
+        return
+
+    print("\n✅ Multi-turn conversation test completed successfully!")
+
+
+async def test_with_workflow_tools():
+    """Test with MassGen workflow tools."""
+
+    print("\n" + "=" * 60)
+    print("🛠️  Testing with workflow tools...")
+
+    backend = ClaudeCodeBackend()
+
+    # Define workflow tools
+    workflow_tools = [
+        {
+            "type": "function",
+            "function": {
+                "name": "new_answer",
+                "description": "Provide an improved answer to the ORIGINAL MESSAGE",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "content": {
+                            "type": "string",
+                            "description": "Your improved answer",
+                        },
+                    },
+                    "required": ["content"],
+                },
+            },
+        },
+        {
+            "type": "function",
+            "function": {
+                "name": "vote",
+                "description": "Vote for the best agent to present final answer",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "agent_id": {
+                            "type": "string",
+                            "description": "Agent ID to vote for",
+                        },
+                        "reason": {
+                            "type": "string",
+                            "description": "Reason for voting",
+                        },
+                    },
+                    "required": ["agent_id", "reason"],
+                },
+            },
+        },
+    ]
+
+    messages = [
+        {
+            "role": "user",
+            "content": "You are participating in a multi-agent workflow. Please provide an answer about the benefits of Python programming, then use the new_answer tool to submit your response.",
+        },
+    ]
+
+    try:
+        print(f"\n📨 Sending workflow tool message with {len(workflow_tools)} tools")
+        print(f"🛠️  Tools: {[t['function']['name'] for t in workflow_tools]}")
+
+        print("\n📡 Stream chunks received:")
+        print("-" * 40)
+
+        chunk_count = 0
+        workflow_content = ""
+        detected_tool_calls = []
+
+        async for chunk in backend.stream_with_tools(messages, workflow_tools):
+            chunk_count += 1
+            print(f"[{chunk_count:2d}] Type: {chunk.type:<20} Source: {chunk.source or 'None':<20}")
+
+            if chunk.type == "content":
+                print(f"     Content: {repr(chunk.content)}")
+                workflow_content += chunk.content or ""
+            elif chunk.type == "tool_calls":
+                print(f"     🛠️  Tool calls detected: {chunk.tool_calls}")
+                detected_tool_calls.extend(chunk.tool_calls or [])
+            elif chunk.type == "complete_response":
+                print(f"     Response metadata: {chunk.response}")
+            elif chunk.type == "done":
+                print("     ✅ Stream completed")
+                break
+
+            print()
+
+        print("-" * 40)
+        print(f"📊 Workflow chunks: {chunk_count}")
+        print(f"📝 Workflow content length: {len(workflow_content)} chars")
+        print(f"🛠️  Detected tool calls: {len(detected_tool_calls)}")
+        for i, tool_call in enumerate(detected_tool_calls):
+            print(f"     [{i+1}] {tool_call.get('function', {}).get('name', 'unknown')}: {tool_call}")
+        print(f"💰 Token usage: {backend.get_token_usage()}")
+
+        if workflow_content:
+            print(f"\n📄 Workflow response:\n{workflow_content}")
+
+    except Exception as e:
+        print(f"❌ Error during workflow test: {e}")
+        import traceback
+
+        traceback.print_exc()
+        return
+
+    print("\n✅ Workflow tools test completed!")
+
+
+async def main():
+    """Run all real tests."""
+    print("🧪 ClaudeCodeBackend Real API Tests")
+    print("=" * 60)
+
+    await test_real_stream_with_tools()
+    await test_with_workflow_tools()
+
+    print("\n🎉 All tests completed!")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())