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

---
 examples/demos/spike/agent_rearrange_test.py | 16 +++++++---------
 1 file changed, 7 insertions(+), 9 deletions(-)

diff --git a/examples/demos/spike/agent_rearrange_test.py b/examples/demos/spike/agent_rearrange_test.py
index e6aa044d..3340091a 100644
--- a/examples/demos/spike/agent_rearrange_test.py
+++ b/examples/demos/spike/agent_rearrange_test.py
@@ -9,7 +9,7 @@ Todo
 import os
 from dotenv import load_dotenv
 from swarms import Agent, AgentRearrange
-from swarm_models import OpenAIChat, OpenAIFunctionCaller
+from swarms.utils.litellm_wrapper import LiteLLM
 from pydantic import BaseModel
 from typing import List
 
@@ -31,10 +31,8 @@ load_dotenv()
 api_key = os.getenv("GROQ_API_KEY")
 
 # Initialize the model
-model = OpenAIChat(
-    openai_api_base="https://api.groq.com/openai/v1",
-    openai_api_key=api_key,
-    model_name="llama-3.1-70b-versatile",
+model = LiteLLM(
+    model_name="groq/llama-3.1-70b-versatile",
     temperature=0.1,
 )
 
@@ -52,11 +50,11 @@ You are a college selection final decision maker. Your role is to:
     
 """
 
-function_caller = OpenAIFunctionCaller(
+function_caller = LiteLLM(
+    model_name="gpt-4o",
     system_prompt=FINAL_AGENT_PROMPT,
-    openai_api_key=os.getenv("OPENAI_API_KEY"),
-    base_model=CollegesRecommendation,
-    parallel_tool_calls=True,
+    response_format=CollegesRecommendation,
+    temperature=0.1,
 )
 
 # Student Profile Analyzer Agent