feat: allow users to pass api_key (#35)

Author: k223kim

src/litai/llm.py

@@ -82,8 +82,7 @@ def __init__(
         fallback_models: Optional[List[str]] = None,
         teamspace: Optional[str] = None,
         max_retries: int = 3,
-        lightning_api_key: Optional[str] = None,
-        lightning_user_id: Optional[str] = None,
+        api_key: Optional[str] = None,
         enable_async: Optional[bool] = False,
         verbose: int = 0,
         full_response: Optional[bool] = None,
@@ -96,26 +95,15 @@ def __init__(
                                                    if the main model fails. Defaults to None.
             teamspace (Optional[List[str]]): Teamspace used for billing.
             max_retries (int): The maximum number of retries for API requests. Defaults to 3.
-            lightning_api_key (Optional[str]): The API key for Lightning AI. Defaults to None.
-            lightning_user_id (Optional[str]): The user ID for Lightning AI. Defaults to None.
+            api_key (Optional[str]): The API key for Lightning AI. Defaults to None.
             enable_async (Optional[bool]): Enable async requests. Defaults to True.
             verbose (int): Verbosity level for logging. Defaults to 0. Must be 0, 1, or 2.
             full_response (bool): Whether the entire response should be returned from the chat
         """
-        if (lightning_api_key is None) != (lightning_user_id is None):
-            missing_param = "lightning_api_key" if lightning_api_key is None else "lightning_user_id"
-            raise ValueError(
-                f"Missing required parameter: '{missing_param}'. "
-                "Both 'lightning_api_key' and 'lightning_user_id' must be provided together. "
-                "Either provide both or none.\n"
-                "To find the API key and user ID, go to the Global Settings page in your Lightning account."
-            )
-
-        if lightning_api_key is not None and lightning_user_id is not None:
-            os.environ["LIGHTNING_API_KEY"] = lightning_api_key
-            os.environ["LIGHTNING_USER_ID"] = lightning_user_id
+        if api_key is not None:
+            os.environ["LIGHTNING_API_KEY"] = api_key
 
-        if os.environ.get("LIGHTNING_API_KEY") is None and os.environ.get("LIGHTNING_USER_ID") is None:
+        if os.environ.get("LIGHTNING_API_KEY") is None:
             self._authenticate()
 
         if verbose not in [0, 1, 2]:

tests/test_llm.py

@@ -21,9 +21,8 @@ def test_initialization_with_config_file(monkeypatch):
     """Test LitAI config."""
     mock_llm_instance = MagicMock()
     monkeypatch.setattr("litai.llm.SDKLLM", mock_llm_instance)
-    LLM(model="openai/gpt-4", lightning_api_key="my-key", lightning_user_id="my-user-id")
+    LLM(model="openai/gpt-4", api_key="my-key")
     assert os.getenv("LIGHTNING_API_KEY") == "my-key"
-    assert os.getenv("LIGHTNING_USER_ID") == "my-user-id"
 
 
 @patch("litai.llm.SDKLLM")
@@ -296,14 +295,13 @@ def mock_auth_constructor():
     monkeypatch.setattr("litai.llm.login.Auth", mock_auth_constructor)
 
     # Test case 1: Both api_key and user_id provided
-    LLM(model="openai/gpt-4", lightning_api_key="my-key", lightning_user_id="my-user-id")
+    LLM(model="openai/gpt-4", api_key="my-key")
 
     # Verify that the authentication was not called
     mock_auth.authenticate.assert_not_called()
 
     # Verify that environment variables were set
     assert os.getenv("LIGHTNING_API_KEY") == "my-key"
-    assert os.getenv("LIGHTNING_USER_ID") == "my-user-id"
 
     # Test case 2: Neither api_key nor user_id provided
     mock_auth.reset_mock()