From a40abbf4900180b721420637b272ddb2c19e9438 Mon Sep 17 00:00:00 2001
From: Lorenze Jay <lorenzejaytech@gmail.com>
Date: Tue, 11 Mar 2025 15:57:08 -0700
Subject: [PATCH] Update AISuiteLLM and LLM utility type handling

- Modify AISuiteLLM to support more flexible input types for messages
- Update type hints in AISuiteLLM to allow string or list of message dictionaries
- Enhance LLM utility function to support broader LLM type annotations
- Remove default `self.stop` attribute from BaseLLM initialization
---
 src/crewai/llms/base_llm.py             | 1 -
 src/crewai/llms/third_party/ai_suite.py | 8 ++++----
 src/crewai/utilities/llm_utils.py       | 4 ++--
 3 files changed, 6 insertions(+), 7 deletions(-)

diff --git a/src/crewai/llms/base_llm.py b/src/crewai/llms/base_llm.py
index 68b7cf1e3..6d832eed7 100644
--- a/src/crewai/llms/base_llm.py
+++ b/src/crewai/llms/base_llm.py
@@ -31,7 +31,6 @@ class BaseLLM(ABC):
         All custom LLM implementations should call super().__init__() to ensure
         that these default attributes are properly initialized.
         """
-        self.stop = []
         self.model = model
         self.temperature = temperature
 
diff --git a/src/crewai/llms/third_party/ai_suite.py b/src/crewai/llms/third_party/ai_suite.py
index f983955d5..ef02f14b3 100644
--- a/src/crewai/llms/third_party/ai_suite.py
+++ b/src/crewai/llms/third_party/ai_suite.py
@@ -1,4 +1,4 @@
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 
 import aisuite as ai
 
@@ -12,11 +12,11 @@ class AISuiteLLM(BaseLLM):
 
     def call(
         self,
-        messages: List[Dict[str, str]],
+        messages: Union[str, List[Dict[str, str]]],
         tools: Optional[List[dict]] = None,
         callbacks: Optional[List[Any]] = None,
         available_functions: Optional[Dict[str, Any]] = None,
-    ) -> str:
+    ) -> Union[str, Any]:
         completion_params = self._prepare_completion_params(messages)
         # print(f"Completion params: {completion_params}")
         response = self.client.chat.completions.create(**completion_params)
@@ -26,7 +26,7 @@ class AISuiteLLM(BaseLLM):
         return response.choices[0].message.content
 
     def _prepare_completion_params(
-        self, messages: List[Dict[str, str]]
+        self, messages: Union[str, List[Dict[str, str]]]
     ) -> Dict[str, Any]:
         print(f"Preparing completion params for {self.model}")
         # print(f"Messages: {messages}")
diff --git a/src/crewai/utilities/llm_utils.py b/src/crewai/utilities/llm_utils.py
index 4cbd4a534..7bd9ebac8 100644
--- a/src/crewai/utilities/llm_utils.py
+++ b/src/crewai/utilities/llm_utils.py
@@ -6,8 +6,8 @@ from crewai.llm import LLM, BaseLLM
 
 
 def create_llm(
-    llm_value: Union[str, BaseLLM, Any, None] = None,
-) -> Optional[BaseLLM]:
+    llm_value: Union[str, BaseLLM, LLM, Any, None] = None,
+) -> Optional[BaseLLM | LLM]:
     """
     Creates or returns an LLM instance based on the given llm_value.