feat: add billing (#47)

* feat: add billing

* nit: update sdk

* fix: test

Author: k223kim

requirements.txt

@@ -1 +1 @@
-lightning_sdk >= 2025.08.08
+lightning_sdk >= 2025.08.08.post1

src/litai/llm.py

@@ -80,7 +80,7 @@ def __init__(
         self,
         model: Optional[ModelLiteral] = None,
         fallback_models: Optional[List[str]] = None,
-        teamspace: Optional[str] = None,
+        billing: Optional[str] = None,
         max_retries: int = 3,
         api_key: Optional[str] = None,
         enable_async: Optional[bool] = False,
@@ -93,7 +93,7 @@ def __init__(
             model (Optional[str]): The main model to use. Defaults to openai/gpt-4o.
             fallback_models (Optional[List[str]]): A list of fallback models to use
                                                    if the main model fails. Defaults to None.
-            teamspace (Optional[List[str]]): Teamspace used for billing.
+            billing (Optional[List[str]]): Teamspace used for billing.
             max_retries (int): The maximum number of retries for API requests. Defaults to 3.
             api_key (Optional[str]): The API key for Lightning AI. Defaults to None.
             enable_async (Optional[bool]): Enable async requests. Defaults to True.
@@ -113,7 +113,7 @@ def __init__(
 
         self._model = model
         self._fallback_models = fallback_models or []
-        self._teamspace = teamspace
+        self._teamspace = billing
         self._enable_async = enable_async
         self._verbose = verbose
         self.max_retries = max_retries

tests/test_llm.py

@@ -48,6 +48,15 @@ def test_default_model(monkeypatch):
         assert llm.model == "openai/gpt-4o"
 
 
+def test_billing(monkeypatch):
+    """Test billing."""
+    mock_llm_instance = MagicMock()
+    monkeypatch.setattr("litai.llm.SDKLLM", mock_llm_instance)
+    llm = LLM(model="openai/gpt-4", billing="my-org/my-teamspace")
+
+    assert llm._teamspace == "my-org/my-teamspace"
+
+
 @patch("litai.llm.SDKLLM")
 def test_cloudy_models_preload(mock_sdkllm):
     """Test that CLOUDY_MODELS are preloaded during LLM initialization."""
@@ -327,7 +336,7 @@ def test_llm_classify_method(mock_sdkllm_class):
     assert result == "positive"
 
     # Test another classification
-    result = llm.classify("this movie was awful.", ["positive", "negative"])
+    result = llm.classify("this movie was really bad.", ["positive", "negative"])
     assert result == "negative"
 
     # Test with multiple classes