From 4442ab5e03f7fffb29600e9570a5585622984d63 Mon Sep 17 00:00:00 2001
From: CI-DEV <154627941+IlumCI@users.noreply.github.com>
Date: Mon, 29 Sep 2025 18:58:05 +0300
Subject: [PATCH] Update test_comprehensive_test.py

---
 tests/test_comprehensive_test.py | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/tests/test_comprehensive_test.py b/tests/test_comprehensive_test.py
index 2ef00e1e..bff0e699 100644
--- a/tests/test_comprehensive_test.py
+++ b/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 (
-            OpenAIFunctionCaller,
-        )
+        from swarms.utils.litellm_wrapper import LiteLLM
         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(
-            base_model=SwarmSpec,
-            api_key=API_KEY,
+        director = LiteLLM(
+            model_name="gpt-4o",
+            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"