mims-harvard__tooluniverse/examples/tool_finder_llm_example.py

#!/usr/bin/env python3
"""
Compatibility test for ToolFinderLLM - ensuring seamless replacement of original tool_finder.py

This script tests that ToolFinderLLM has the same interface and behavior as the original
ToolFinderEmbedding, making it a drop-in replacement.
"""

import sys
import traceback
import warnings

# Suppress RDKit warnings and pkg_resources warnings
warnings.filterwarnings("ignore", category=RuntimeWarning, module="importlib._bootstrap")
warnings.filterwarnings("ignore", message=".*RDKit.*")
warnings.filterwarnings("ignore", message=".*pkg_resources.*")


def test_compatibility():
    """Test that ToolFinderLLM matches original tool_finder interface."""
    print("🧪 Testing ToolFinderLLM compatibility...")

    try:
        from tooluniverse.execute_function import ToolUniverse
        from tooluniverse.tool_finder_llm import ToolFinderLLM

        # Initialize ToolUniverse
        tooluniverse = ToolUniverse()
        tooluniverse.load_tools()

        # Configure LLM tool finder
        config = {
            "name": "ToolFinderLLM",
            "type": "ToolFinderLLM",
            "configs": {
                "api_type": "CHATGPT",
                "model_id": "gpt-4o-1120",
                "temperature": 0.1,
                "max_new_tokens": 2048,
                "return_json": True,
            },
        }

        llm_finder = ToolFinderLLM(config, tooluniverse)

        # Test 1: Basic find_tools() method
        print("✓ Testing find_tools() method...")
        result1 = llm_finder.find_tools(
            message="Find tools for drug safety analysis", rag_num=3
        )
        assert isinstance(result1, list), f"Expected list, got {type(result1)}"
        print(f"  Result type: {type(result1)} ✓")

        # Test 2: find_tools() with return_call_result=True
        print("✓ Testing find_tools() with return_call_result=True...")
        result2 = llm_finder.find_tools(
            message="Find tools for gene analysis", rag_num=2, return_call_result=True
        )
        assert isinstance(result2, tuple), f"Expected tuple, got {type(result2)}"
        assert len(result2) == 2, f"Expected tuple of length 2, got {len(result2)}"
        prompts, tool_names = result2
        assert isinstance(
            prompts, list
        ), f"Expected prompts as list, got {type(prompts)}"
        assert isinstance(
            tool_names, list
        ), f"Expected tool_names as list, got {type(tool_names)}"
        print(
            f"  Result: ({type(prompts)}, {type(tool_names)}) with {len(tool_names)} tools ✓"
        )

        # Test 3: run() method compatibility
        print("✓ Testing run() method...")
        arguments = {
            "description": "Find tools for protein analysis",
            "limit": 2,
            "return_call_result": False,
        }
        result3 = llm_finder.run(arguments)
        # Note: run() method returns string, not list (different from find_tools())
        assert isinstance(
            result3, str
        ), f"Expected string from run(), got {type(result3)}"
        print(f"  Run method result type: {type(result3)} ✓")

        # Test 4: run() with return_call_result=True
        print("✓ Testing run() with return_call_result=True...")
        arguments_with_result = {
            "description": "Find tools for clinical trials",
            "limit": 2,
            "return_call_result": True,
        }
        result4 = llm_finder.run(arguments_with_result)
        assert isinstance(
            result4, tuple
        ), f"Expected tuple from run(), got {type(result4)}"
        print(f"  Run method with return_call_result: {type(result4)} ✓")

        # Test 5: test with picked_tool_names (original interface)
        print("✓ Testing find_tools() with picked_tool_names...")
        test_tools = ["EuropePMC_search_articles", "PubChem_get_CID_by_compound_name"]
        result5 = llm_finder.find_tools(
            picked_tool_names=test_tools, return_call_result=False
        )
        assert isinstance(
            result5, list
        ), f"Expected list with picked_tool_names, got {type(result5)}"
        print(f"  Picked tools result type: {type(result5)} ✓")

        # Test 6: test picked_tool_names with return_call_result=True
        print(
            "✓ Testing find_tools() with picked_tool_names and return_call_result=True..."
        )
        result6 = llm_finder.find_tools(
            picked_tool_names=test_tools, return_call_result=True
        )
        assert isinstance(
            result6, tuple
        ), f"Expected tuple with picked_tool_names, got {type(result6)}"
        prompts6, names6 = result6
        assert isinstance(
            prompts6, list
        ), f"Expected prompts as list, got {type(prompts6)}"
        assert isinstance(names6, list), f"Expected names as list, got {type(names6)}"
        assert len(names6) <= len(
            test_tools
        ), f"Expected names count <= input, got {len(names6)} > {len(test_tools)}"
        print(
            f"  Picked tools with return_call_result: ({type(prompts6)}, {type(names6)}) ✓"
        )

        print("\n🎉 All compatibility tests passed!")
        print("✅ ToolFinderLLM interface matches original ToolFinderEmbedding")
        return True

    except ImportError as e:
        print(f"⚠️ ToolUniverse not available: {e}")
        print("Skipping tests - this is expected in some environments")
        return True

    except Exception as e:
        print(f"❌ Compatibility test failed: {e}")
        traceback.print_exc()
        return False


def test_error_handling():
    """Test error handling and edge cases."""
    print("\n🧪 Testing error handling...")

    try:
        from tooluniverse.execute_function import ToolUniverse
        from tooluniverse.tool_finder_llm import ToolFinderLLM

        # Initialize ToolUniverse
        tooluniverse = ToolUniverse()
        tooluniverse.load_tools()

        config = {
            "name": "ToolFinderLLM",
            "type": "ToolFinderLLM",
            "configs": {
                "api_type": "CHATGPT",
                "model_id": "gpt-4o-1120",
                "temperature": 0.1,
                "max_new_tokens": 2048,
                "return_json": True,
            },
        }

        llm_finder = ToolFinderLLM(config, tooluniverse)

        # Test with invalid arguments
        print("✓ Testing error handling with None arguments...")
        try:
            # This should raise an assertion error
            result = llm_finder.find_tools(message=None, picked_tool_names=None)
            print(f"  ❌ Expected AssertionError but got result: {type(result)}")
            return False
        except AssertionError:
            print("  ✓ AssertionError raised as expected")

        # Test with empty picked_tool_names
        print("✓ Testing with empty picked_tool_names...")
        result = llm_finder.find_tools(picked_tool_names=[], return_call_result=False)
        assert isinstance(
            result, list
        ), f"Expected list with empty picked_tool_names, got {type(result)}"
        assert len(result) == 0, f"Expected empty result, got {len(result)} items"
        print("  ✓ Empty picked_tool_names handled correctly")

        # Test run() with missing required arguments
        print("✓ Testing run() with empty arguments...")
        try:
            result = llm_finder.run({})
            print(f"  ❌ Expected AssertionError but got result: {type(result)}")
            return False
        except AssertionError:
            print("  ✓ AssertionError raised as expected for empty arguments")

        print("✅ Error handling tests passed!")
        return True

    except ImportError as e:
        print(f"⚠️ ToolUniverse not available: {e}")
        return True

    except Exception as e:
        print(f"❌ Error handling test failed: {e}")
        traceback.print_exc()
        return False


def main():
    """Run compatibility tests to ensure seamless replacement of original tool_finder."""
    print("🚀 ToolFinderLLM Compatibility Tests")
    print("=" * 60)

    tests = [
        ("Interface Compatibility", test_compatibility),
        ("Error Handling", test_error_handling),
    ]

    results = []
    for test_name, test_func in tests:
        try:
            success = test_func()
            results.append((test_name, success))
        except Exception as e:
            print(f"❌ {test_name} crashed: {e}")
            results.append((test_name, False))

    print("\n" + "=" * 60)
    print("📊 Test Results Summary:")
    print("=" * 60)

    passed = 0
    for test_name, success in results:
        status = "✅ PASS" if success else "❌ FAIL"
        print(f"{status} | {test_name}")
        if success:
            passed += 1

    print(f"\n🏁 Tests completed: {passed}/{len(results)} passed")

    if passed == len(results):
        print("🎉 All compatibility tests passed!")
        print(
            "✅ ToolFinderLLM can seamlessly replace the original ToolFinderEmbedding"
        )
        return 0
    else:
        print("⚠️  Some compatibility tests failed.")
        print("❌ Interface modifications needed before replacement")
        return 1


if __name__ == "__main__":
    sys.exit(main())