From 6497e5d82779fd47301d5e5f2891422757890f47 Mon Sep 17 00:00:00 2001
From: CI-DEV <154627941+IlumCI@users.noreply.github.com>
Date: Thu, 14 Aug 2025 23:58:33 +0300
Subject: [PATCH] Update litellm_wrapper.py

---
 swarms/utils/litellm_wrapper.py | 32 +++++++++++++++++++++++++-------
 1 file changed, 25 insertions(+), 7 deletions(-)

diff --git a/swarms/utils/litellm_wrapper.py b/swarms/utils/litellm_wrapper.py
index 1b7d3c60..e438a3e3 100644
--- a/swarms/utils/litellm_wrapper.py
+++ b/swarms/utils/litellm_wrapper.py
@@ -214,14 +214,32 @@ class LiteLLM:
 
     def output_for_tools(self, response: any):
         if self.mcp_call is True:
-            out = response.choices[0].message.tool_calls[0].function
-            output = {
-                "function": {
-                    "name": out.name,
-                    "arguments": out.arguments,
+            # Check if there are any tool calls in the response
+            if (hasattr(response, 'choices') and 
+                response.choices and 
+                hasattr(response.choices[0], 'message') and
+                hasattr(response.choices[0].message, 'tool_calls') and
+                response.choices[0].message.tool_calls):
+                
+                # Extract the first tool call
+                out = response.choices[0].message.tool_calls[0].function
+                output = {
+                    "function": {
+                        "name": out.name,
+                        "arguments": out.arguments,
+                    }
                 }
-            }
-            return output
+                return output
+            else:
+                # No tool calls present, return the regular content
+                if (hasattr(response, 'choices') and 
+                    response.choices and 
+                    hasattr(response.choices[0], 'message') and
+                    hasattr(response.choices[0].message, 'content')):
+                    return response.choices[0].message.content
+                else:
+                    # Fallback: return the response as is
+                    return response
         else:
             out = response.choices[0].message.tool_calls