feat - improve classify to use list (#14)

* comments

* added tests

* fix prompt

* fix tests

* fix tests

* lint

* make input an array

Author: williamFalcon

src/litai/client.py

@@ -452,27 +452,28 @@ def if_(self, input: str, question: str) -> bool:
         response = self.chat(prompt).strip().lower()
         return "yes" in response
 
-    def classify(self, input: str, *choices: str) -> str:
-        """Returns the label the model chooses from the given options.
+    def classify(self, input: str, choices: List[str]) -> str:
+        """Returns the label the model chooses from the given list of options.
 
         Example:
-            llm.classify("This product sucks.", "positive", "negative") → "negative"
+            llm.classify("This product sucks.", ["positive", "negative", "neutral"]) → "negative"
         """
         normalized_choices = [c.strip().lower() for c in choices]
         choices_str = ", ".join(normalized_choices)
+
         prompt = f"""
-        You are given this input
+        You are given this input:
         <input>
-        {input}
+        {input.strip()}
         </input>
 
         And the following choices:
         <choices>
         {choices_str}
         </choices>
 
-        Answer with only one of the choices
-        """
+        Answer with only one of the choices.
+        """.strip()
 
         response = self.chat(prompt).strip().lower()
 

tests/test_llm.py

@@ -329,13 +329,13 @@ def test_llm_classify_method(mock_sdkllm_class):
     mock_sdkllm_instance.chat.side_effect = ["positive", "negative", "neutral"]
 
     # Test simple classification
-    result = llm.classify("this movie was great!", "positive", "negative")
+    result = llm.classify("this movie was great!", ["positive", "negative"])
     assert result == "positive"
 
     # Test another classification
-    result = llm.classify("this movie was awful.", "positive", "negative")
+    result = llm.classify("this movie was awful.", ["positive", "negative"])
     assert result == "negative"
 
     # Test with multiple classes
-    result = llm.classify("it was okay.", "positive", "negative", "neutral")
+    result = llm.classify("it was okay.", ["positive", "negative", "neutral"])
     assert result == "neutral"