massgen 0.1.0a3__py3-none-any.whl → 0.1.2__py3-none-any.whl

massgen/tests/test_tools.py

@@ -0,0 +1,127 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""Test script for MassGen tool implementation."""
+
+import asyncio
+
+from massgen.tool._code_executors import run_python_script
+from massgen.tool._file_handlers import read_file_content
+from massgen.tool._manager import ToolManager
+from massgen.tool._result import ExecutionResult, TextContent
+
+
+async def sample_math_tool(x: int, y: int) -> ExecutionResult:
+    """Add two numbers together.
+
+    Args:
+        x: First number
+        y: Second number
+
+    Returns:
+        Sum of the two numbers
+    """
+    result = x + y
+    return ExecutionResult(
+        output_blocks=[
+            TextContent(data=f"The sum of {x} and {y} is {result}"),
+        ],
+    )
+
+
+async def test_tool_manager():
+    """Test the tool manager functionality."""
+    print("Testing MassGen Tool Manager\n" + "=" * 40)
+
+    # Create manager
+    manager = ToolManager()
+    print("✓ Tool manager created")
+
+    # Create categories
+    manager.setup_category(
+        category_name="math",
+        description="Mathematical operations",
+        enabled=True,
+        usage_hints="Use these tools for calculations",
+    )
+    print("✓ Created 'math' category")
+
+    manager.setup_category(
+        category_name="file_ops",
+        description="File operations",
+        enabled=False,
+    )
+    print("✓ Created 'file_ops' category")
+
+    # Register tools
+    manager.add_tool_function(
+        func=sample_math_tool,
+        category="math",
+        description="Adds two numbers",
+    )
+    print("✓ Registered sample_math_tool")
+
+    manager.add_tool_function(
+        func=run_python_script,
+        category="default",
+        description="Execute Python code",
+    )
+    print("✓ Registered run_python_script")
+
+    manager.add_tool_function(
+        func=read_file_content,
+        category="file_ops",
+        description="Read file contents",
+    )
+    print("✓ Registered read_file_content")
+
+    # Get schemas
+    schemas = manager.fetch_tool_schemas()
+    print(f"\n✓ Active tool schemas: {len(schemas)} tools")
+    for schema in schemas:
+        print(f"  - {schema['function']['name']}: {schema['function'].get('description', 'No description')}")
+
+    # Test tool execution
+    print("\nTesting tool execution:")
+
+    # Test math tool
+    tool_request = {
+        "name": "sample_math_tool",
+        "input": {"x": 5, "y": 3},
+    }
+
+    async for result in manager.execute_tool(tool_request):
+        print(f"  Math result: {result.output_blocks[0].data}")
+
+    # Test Python execution
+    python_request = {
+        "name": "run_python_script",
+        "input": {
+            "source_code": "print('Hello from MassGen!')\nprint(2 + 2)",
+        },
+    }
+
+    async for result in manager.execute_tool(python_request):
+        output = result.output_blocks[0].data
+        if "<stdout>" in output:
+            stdout_start = output.find("<stdout>") + 8
+            stdout_end = output.find("</stdout>")
+            stdout = output[stdout_start:stdout_end]
+            print(f"  Python output: {stdout.strip()}")
+
+    # Test enabling file_ops category
+    print("\nEnabling file_ops category...")
+    manager.modify_categories(["file_ops"], enabled=True)
+    schemas_after = manager.fetch_tool_schemas()
+    print(f"✓ Active tools after enabling: {len(schemas_after)}")
+
+    # Get category hints
+    hints = manager.fetch_category_hints()
+    if hints:
+        print(f"\nCategory hints:\n{hints}")
+
+    print("\n" + "=" * 40)
+    print("All tests completed successfully! ✓")
+
+
+if __name__ == "__main__":
+    asyncio.run(test_tool_manager())

massgen/token_manager/token_manager.py

@@ -59,10 +59,19 @@ class TokenCostCalculator:
             "o3-mini": ModelPricing(0.0011, 0.0044, 200000, 100000),
         },
         "Anthropic": {
-            "claude-3-5-sonnet": ModelPricing(0.003, 0.015, 200000, 8192),
-            "claude-3-5-haiku": ModelPricing(0.001, 0.005, 200000, 8192),
-            "claude-3-opus": ModelPricing(0.015, 0.075, 200000, 4096),
-            "claude-3-sonnet": ModelPricing(0.003, 0.015, 200000, 4096),
+            # Claude 4.5 models (October 2024+)
+            "claude-haiku-4-5": ModelPricing(0.001, 0.005, 200000, 65536),  # $1/MTok input, $5/MTok output, 64K max output
+            "claude-sonnet-4-5": ModelPricing(0.003, 0.015, 200000, 65536),  # $3/MTok input, $15/MTok output, 64K max output
+            # Claude 4 models
+            "claude-opus-4.1": ModelPricing(0.015, 0.075, 200000, 32768),  # $15/MTok input, $75/MTok output, 32K max output
+            "claude-opus-4": ModelPricing(0.015, 0.075, 200000, 32768),  # $15/MTok input, $75/MTok output, 32K max output
+            "claude-sonnet-4": ModelPricing(0.003, 0.015, 200000, 8192),  # $3/MTok input, $15/MTok output
+            # Claude 3.5 models
+            "claude-3-5-sonnet": ModelPricing(0.003, 0.015, 200000, 8192),  # $3/MTok input, $15/MTok output
+            "claude-3-5-haiku": ModelPricing(0.0008, 0.004, 200000, 8192),  # $0.80/MTok input, $4/MTok output
+            # Claude 3 models (deprecated)
+            "claude-3-opus": ModelPricing(0.015, 0.075, 200000, 4096),  # Deprecated
+            "claude-3-sonnet": ModelPricing(0.003, 0.015, 200000, 4096),  # Deprecated
             "claude-3-haiku": ModelPricing(0.00025, 0.00125, 200000, 4096),
         },
         "Google": {