diff --git a/src/litai/llm.py b/src/litai/llm.py index 0784dcd..4bc6bc9 100644 --- a/src/litai/llm.py +++ b/src/litai/llm.py @@ -225,6 +225,7 @@ def _model_call( # noqa: D417 lit_tools: Optional[List[LitTool]] = None, auto_call_tools: bool = False, reasoning_effort: Optional[str] = None, + temperature: Optional[float] = None, **kwargs: Any, ) -> str: """Handles the model call and logs appropriate messages.""" @@ -245,6 +246,7 @@ def _model_call( # noqa: D417 full_response=full_response, tools=tools, reasoning_effort=reasoning_effort, + temperature=temperature, **kwargs, ) if tools and isinstance(response, V1ConversationResponseChunk): @@ -305,6 +307,7 @@ async def async_chat( lit_tools: Optional[List[LitTool]] = None, auto_call_tools: bool = False, reasoning_effort: Optional[str] = None, + temperature: Optional[float] = None, **kwargs: Any, ) -> Union[str, AsyncIterator[str], None]: """Sends a message to the LLM asynchronously with full retry/fallback logic.""" @@ -325,6 +328,7 @@ async def async_chat( full_response=full_response, auto_call_tools=auto_call_tools, reasoning_effort=reasoning_effort, + temperature=temperature, **kwargs, ) @@ -357,6 +361,7 @@ def sync_chat( lit_tools: Optional[List[LitTool]] = None, auto_call_tools: bool = False, reasoning_effort: Optional[str] = None, + temperature: Optional[float] = None, **kwargs: Any, ) -> Union[str, Iterator[str], None]: """Sends a message to the LLM synchronously with full retry/fallback logic.""" @@ -377,6 +382,7 @@ def sync_chat( full_response=full_response, auto_call_tools=auto_call_tools, reasoning_effort=reasoning_effort, + temperature=temperature, **kwargs, ) @@ -416,6 +422,7 @@ def chat( # noqa: D417 tools: Optional[Sequence[Union[LitTool, "StructuredTool"]]] = None, auto_call_tools: bool = False, reasoning_effort: Optional[Literal["none", "low", "medium", "high"]] = None, + temperature: Optional[float] = None, **kwargs: Any, ) -> Union[str, Task[Union[str, AsyncIterator[str], None]], Iterator[str], None]: """Sends a message to the LLM and retrieves a response. @@ -437,6 +444,8 @@ def chat( # noqa: D417 full_response (bool): Whether the entire response should be returned from the chat. auto_call_tools (bool): Tools will be executed automatically whenever applicable. Defaults to False. reasoning_effort (Optional[Literal["low", "medium", "high"]]): The level of reasoning effort for the model. + temperature (Optional[float]): Sampling temperature between 0 and 2. Higher values make output more random, + lower values make it more deterministic. Defaults to None (uses model default). **kwargs (Any): Additional keyword arguments Returns: @@ -445,6 +454,9 @@ def chat( # noqa: D417 if reasoning_effort is not None and reasoning_effort not in ["none", "low", "medium", "high"]: raise ValueError("reasoning_effort must be 'low', 'medium', 'high', or None") + if temperature is not None and not (0.0 <= temperature <= 2.0): + raise ValueError("temperature must be between 0 and 2") + self._wait_for_model() lit_tools = LitTool.convert_tools(tools) processed_tools = [tool.as_tool() for tool in lit_tools] if lit_tools else None @@ -482,6 +494,7 @@ def chat( # noqa: D417 lit_tools=lit_tools, auto_call_tools=auto_call_tools, reasoning_effort=reasoning_effort, + temperature=temperature, **kwargs, ) ) @@ -501,6 +514,7 @@ def chat( # noqa: D417 lit_tools=lit_tools, auto_call_tools=auto_call_tools, reasoning_effort=reasoning_effort, + temperature=temperature, **kwargs, ) diff --git a/tests/test_llm.py b/tests/test_llm.py index 233aa89..d784cbd 100644 --- a/tests/test_llm.py +++ b/tests/test_llm.py @@ -136,6 +136,7 @@ def test_llm_chat(mock_llm_class): my_kwarg="test-kwarg", tools=None, reasoning_effort=None, + temperature=None, ) test_kwargs = mock_llm_instance.chat.call_args.kwargs assert test_kwargs.get("my_kwarg") == "test-kwarg" @@ -144,6 +145,37 @@ def test_llm_chat(mock_llm_class): mock_llm_instance.reset_conversation.assert_called_once() +@patch("litai.llm.SDKLLM") +def test_llm_chat_temperature(mock_llm_class): + """Test that temperature is passed through to the underlying model.""" + from litai.llm import LLM as LLMCLIENT + + LLMCLIENT._sdkllm_cache.clear() + mock_llm_instance = MagicMock() + mock_llm_instance.chat.return_value = "Creative response." + mock_llm_class.return_value = mock_llm_instance + + llm = LLM(model="openai/gpt-4") + response = llm.chat("Tell me a story.", temperature=0.9) + + assert response == "Creative response." + call_kwargs = mock_llm_instance.chat.call_args.kwargs + assert call_kwargs["temperature"] == 0.9 + + +@patch("litai.llm.SDKLLM") +def test_llm_chat_temperature_validation(mock_llm_class): + """Test that invalid temperature values raise ValueError.""" + mock_llm_class.return_value = MagicMock() + llm = LLM(model="openai/gpt-4") + + with pytest.raises(ValueError, match="temperature must be between 0 and 2"): + llm.chat("Hello", temperature=3.0) + + with pytest.raises(ValueError, match="temperature must be between 0 and 2"): + llm.chat("Hello", temperature=-0.1) + + def test_model_override(monkeypatch): """Test override model logic when main model fails.""" mock_llm = MagicMock() @@ -193,6 +225,7 @@ def mock_llm_constructor(name, teamspace="default-teamspace", **kwargs): full_response=True, tools=None, reasoning_effort=None, + temperature=None, ) @@ -244,6 +277,7 @@ def mock_llm_constructor(name, teamspace="default-teamspace", **kwargs): full_response=False, tools=None, reasoning_effort=None, + temperature=None, ) @@ -295,6 +329,7 @@ def mock_llm_constructor(name, teamspace="default-teamspace", **kwargs): full_response=False, tools=None, reasoning_effort=None, + temperature=None, )