Add temperature parameter to LLM.chat()

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,
         )
 

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,
     )