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
This commit is contained in:
Lorenze Jay
2025-03-11 15:57:08 -07:00
parent 25c64ae86d
commit a40abbf490
3 changed files with 6 additions and 7 deletions

View File

@@ -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

View File

@@ -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}")

View File

@@ -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.