fix CI wrong assert

src/litai/client.py

@@ -426,9 +426,7 @@ def list_conversations(self) -> List[str]:
         return self._llm.list_conversations()
 
     def if_(self, input: str, choice1: Optional[str] = None, choice2: Optional[str] = None) -> bool:
-        """Returns True if the model selects the first choice, False otherwise.
-        Defaults to a yes/no binary decision.
-        """
+        """Returns True if the model selects the first choice, False otherwise. Defaults to a yes/no binary decision."""
         choice1 = (choice1 or "yes").strip().lower()
         choice2 = (choice2 or "no").strip().lower()
 
@@ -438,11 +436,10 @@ def if_(self, input: str, choice1: Optional[str] = None, choice2: Optional[str]
 
         if response == choice1:
             return True
-        elif response == choice2:
+        if response == choice2:  # noqa: SIM103
             return False
-        else:
-            # fallback: assume choice1 if unclear
-            return True
+        # fallback: assume choice1 if unclear
+        return True
 
     def classify(self, input: str, *choices: str) -> str:
         """Returns the label the model chooses from the given options.

tests/test_llm.py

@@ -15,7 +15,7 @@ def test_initialization_with_config_file(monkeypatch):
     monkeypatch.setattr("litai.client.SDKLLM", mock_llm_instance)
     LLM(model="openai/gpt-4", lightning_api_key="my-key", lightning_user_id="my-user-id")
     assert os.getenv("LIGHTNING_API_KEY") == "my-key"
-    assert os.getenv("LIGHTNING_USER_ID") == "my-user_id"
+    assert os.getenv("LIGHTNING_USER_ID") == "my-user-id"
 
 
 @patch("litai.client.SDKLLM")
@@ -293,9 +293,10 @@ def mock_auth_constructor():
 def test_llm_if_method(mock_llm_class):
     """Test the LLM if_ method."""
     from litai.client import LLM as LLMCLIENT
+
     LLMCLIENT._sdkllm_cache.clear()
     mock_llm_instance = MagicMock()
-    
+
     # Test case where the condition is true
     mock_llm_instance.chat.return_value = "yes"
     mock_llm_class.return_value = mock_llm_instance
@@ -311,7 +312,7 @@ def test_llm_if_method(mock_llm_class):
         stream=False,
         full_response=False,
     )
-    
+
     # Test case where the condition is false
     mock_llm_instance.chat.return_value = "no"
     assert llm.if_("is it false?") is False
@@ -325,26 +326,31 @@ def test_llm_if_method(mock_llm_class):
         stream=False,
         full_response=False,
     )
-    
+
     # Test case with different capitalization/spacing
     mock_llm_instance.chat.return_value = " Yes "
     assert llm.if_("is it a positive response?") is True
 
+
 @patch("litai.client.SDKLLM")
 def test_llm_classify_method(mock_llm_class):
     """Test the LLM classify method."""
     from litai.client import LLM as LLMCLIENT
+
     LLMCLIENT._sdkllm_cache.clear()
     mock_llm_instance = MagicMock()
-    
+
     # Test a simple classification
     mock_llm_instance.chat.return_value = "positive"
     mock_llm_class.return_value = mock_llm_instance
     llm = LLM(model="openai/gpt-4")
     result = llm.classify("this movie was great!", "positive", "negative")
     assert result == "positive"
     mock_llm_instance.chat.assert_called_with(
-        prompt="this movie was great!\n\nclassify the input as one of these: positive, negative. reply with only the class.",
+        prompt=(
+            "this movie was great!\n\nclassify the input as one of these: "
+            "positive, negative. reply with only the class."
+        ),
         system_prompt=None,
         max_completion_tokens=500,
         images=None,
@@ -359,7 +365,10 @@ def test_llm_classify_method(mock_llm_class):
     result = llm.classify("this movie was awful.", "positive", "negative")
     assert result == "negative"
     mock_llm_instance.chat.assert_called_with(
-        prompt="this movie was awful.\n\nclassify the input as one of these: positive, negative. reply with only the class.",
+        prompt=(
+            "this movie was awful.\n\nclassify the input as one of these: "
+            "positive, negative. reply with only the class."
+        ),
         system_prompt=None,
         max_completion_tokens=500,
         images=None,
@@ -374,12 +383,15 @@ def test_llm_classify_method(mock_llm_class):
     result = llm.classify("it was okay.", "positive", "negative", "neutral")
     assert result == "neutral"
     mock_llm_instance.chat.assert_called_with(
-        prompt="it was okay.\n\nclassify the input as one of these: positive, negative, neutral. reply with only the class.",
+        prompt=(
+            "it was okay.\n\nclassify the input as one of these: positive,"
+            " negative, neutral. reply with only the class."
+        ),
         system_prompt=None,
         max_completion_tokens=500,
         images=None,
         conversation=None,
         metadata=None,
         stream=False,
         full_response=False,
-    )
+    )