Update test_comprehensive_test.py

tests/test_comprehensive_test.py

@@ -463,9 +463,7 @@ def test_spreadsheet_swarm():
 def test_hierarchical_swarm():
     """Test HierarchicalSwarm structure"""
     try:
-        from swarms.utils.function_caller_model import (
+        from swarms.utils.litellm_wrapper import LiteLLM
-            OpenAIFunctionCaller,
-        )
         from swarms.structs.hiearchical_swarm import SwarmSpec
 
         # Create worker agents
@@ -481,9 +479,9 @@ def test_hierarchical_swarm():
         ]
 
         # Create director agent with explicit knowledge of available agents
-        director = OpenAIFunctionCaller(
+        director = LiteLLM(
-            base_model=SwarmSpec,
+            model_name="gpt-4o",
-            api_key=API_KEY,
+            response_format=SwarmSpec,
             system_prompt=(
                 "As the Director of this Hierarchical Agent Swarm, you coordinate tasks among agents. "
                 "You must ONLY assign tasks to the following available agents:\n"