Skip to content

Commit bef00d1

Browse files
authored
fix: correct none reasoning mapping for small gemini models (#65)
* allow to pass none as reasoning effort * bump sdk version * fix: None maps to none on small gemini models * fix: None maps to none on small gemini models * fix rebase * fix test
1 parent a370729 commit bef00d1

File tree

2 files changed

+53
-3
lines changed

2 files changed

+53
-3
lines changed

src/litai/llm.py

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,9 @@
4141
"google/gemini-2.5-flash",
4242
}
4343

44+
# this nmodels support reasoning_effort='none'
45+
NONE_REASONING_MODELS = ["google/gemini-2.0-flash", "google/gemini-2.5-flash-lite-preview-06-17"]
46+
4447
logger = logging.getLogger(__name__)
4548

4649

@@ -290,15 +293,19 @@ def chat( # noqa: D417
290293
categorized by conversation ID.
291294
full_response (bool): Whether the entire response should be returned from the chat.
292295
auto_call_tools (bool): Tools will be executed automatically whenever applicable. Defaults to False.
293-
reasoning_effort (Optional[Literal["none", "low", "medium", "high"]]):
294-
The level of reasoning effort for the model.
296+
reasoning_effort (Optional[Literal["low", "medium", "high"]]): The level of reasoning effort for the model.
295297
**kwargs (Any): Additional keyword arguments
296298
297299
Returns:
298300
str: The response from the LLM.
299301
"""
300302
if reasoning_effort is not None and reasoning_effort not in ["none", "low", "medium", "high"]:
301303
raise ValueError("reasoning_effort must be 'low', 'medium', 'high', or None")
304+
if reasoning_effort is None and (
305+
model in NONE_REASONING_MODELS or (self._model in NONE_REASONING_MODELS and model is None)
306+
):
307+
reasoning_effort = "none"
308+
302309
self._wait_for_model()
303310
lit_tools = LitTool.convert_tools(tools)
304311
processed_tools = [tool.as_tool() for tool in lit_tools] if lit_tools else None

tests/test_llm.py

Lines changed: 44 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -119,7 +119,50 @@ def test_llm_chat(mock_llm_class):
119119
system_prompt="You are a helpful assistant.",
120120
metadata={"user_api": "123456"},
121121
my_kwarg="test-kwarg",
122-
reasoning_effort="none",
122+
reasoning_effort=None,
123+
)
124+
125+
assert isinstance(response, str)
126+
assert "helpful" in response.lower()
127+
mock_llm_instance.chat.assert_called_once_with(
128+
prompt="Hello, who are you?",
129+
system_prompt="You are a helpful assistant.",
130+
max_completion_tokens=None,
131+
images=None,
132+
conversation=None,
133+
metadata={"user_api": "123456"},
134+
stream=False,
135+
full_response=False,
136+
my_kwarg="test-kwarg",
137+
tools=None,
138+
reasoning_effort=None,
139+
)
140+
test_kwargs = mock_llm_instance.chat.call_args.kwargs
141+
assert test_kwargs.get("my_kwarg") == "test-kwarg"
142+
143+
llm.reset_conversation("test")
144+
mock_llm_instance.reset_conversation.assert_called_once()
145+
146+
147+
@patch("litai.llm.SDKLLM")
148+
def test_reasoning_effort_override(mock_llm_class):
149+
"""Test LigtningLLM chat."""
150+
from litai.llm import LLM as LLMCLIENT
151+
152+
LLMCLIENT._sdkllm_cache.clear()
153+
mock_llm_instance = MagicMock()
154+
mock_llm_instance.chat.return_value = "Hello! I am a helpful assistant."
155+
156+
mock_llm_class.return_value = mock_llm_instance
157+
158+
llm = LLM(model="google/gemini-2.0-flash")
159+
160+
response = llm.chat(
161+
"Hello, who are you?",
162+
system_prompt="You are a helpful assistant.",
163+
metadata={"user_api": "123456"},
164+
my_kwarg="test-kwarg",
165+
reasoning_effort=None,
123166
)
124167

125168
assert isinstance(response, str)

0 commit comments

Comments
 (0)