Add authentication for local (#3)

* Add teamspace configuration and authentication method to LLM class

* Refactor LLM class by removing teamspace validation and unused authentication call; add unit tests for authentication method

* fix pre-commit

* Update README for clarity on Lightning Studio usage and local authentication; modify LLM class authentication method signature to accept optional parameters.

* update

Author: aniketmaurya

README.md

@@ -50,7 +50,7 @@ pip install litai
 
 ## Run on a Studio
 
-When running inside Lightning Studio, authentication is handled automatically. You can use any available LLM for text generation.
+When running inside Lightning Studio, you can use any available LLM out of the box — no extra setup required.
 
 ```python
 from litai import LLM
@@ -62,22 +62,14 @@ print(llm.chat("who are you?"))
 
 ## Run locally (outside Studio)
 
-To use LitAI outside of Lightning Studio, you’ll need to manually authenticate using the following:
+To use LitAI outside of Lightning Studio, you'll need to explicitly provide your teamspace name.
 
-- **teamspace** — Format: "owner-name/teamspace-name"
-  (e.g. "username/my-team" or "org-name/my-team")
-- **lightning_api_key** — Your personal API key from Global Settings → Keys
-- **lightning_user_id** — Your Lightning user ID, also found in the Keys section
+The teamspace input format is: `"owner-name/teamspace-name"` (e.g. `"username/my-team"` or `"org-name/team-name"`)
 
 ```python
 from litai import LLM
 
-llm = LLM(
-    model="openai/gpt-4",
-    teamspace="owner-name/teamspace-name",
-    lightning_api_key="your-api-key",
-    lightning_user_id="your-user-id",
-)
+llm = LLM(model="openai/gpt-4", teamspace="owner-name/teamspace-name")
 print(llm.chat("who are you?"))
 # I'm an AI by OpenAI
 ```

src/litai/client.py

@@ -22,6 +22,7 @@
 from typing import Any, Dict, List, Optional, Union
 
 import requests
+from lightning_sdk.lightning_cloud import login
 from lightning_sdk.llm import LLM as SDKLLM
 
 from litai.tools import LitTool
@@ -111,6 +112,11 @@ def __init__(
         if lightning_user_id is not None:
             os.environ["LIGHTNING_USER_ID"] = lightning_user_id
 
+        if teamspace is None:
+            teamspace = os.environ.get("LIGHTNING_TEAMSPACE")
+
+        self._authenticate(lightning_api_key, lightning_user_id)
+
         if verbose not in [0, 1, 2]:
             raise ValueError("Verbose must be 0, 1, or 2.")
         self._verbose = verbose
@@ -134,6 +140,19 @@ def __init__(
 
         threading.Thread(target=self._load_models, daemon=True).start()
 
+    def _authenticate(self, lightning_api_key: Optional[str], lightning_user_id: Optional[str]) -> None:
+        if not (lightning_api_key and lightning_user_id):
+            return
+        auth = login.Auth()
+        try:
+            auth.authenticate()
+            user_api_key = auth.api_key
+            user_id = auth.user_id
+            os.environ["LIGHTNING_API_KEY"] = user_api_key
+            os.environ["LIGHTNING_USER_ID"] = user_id
+        except ConnectionError as e:
+            raise e
+
     @property
     def model(self) -> str:
         """Returns the main model name."""

tests/test_llm.py

@@ -255,3 +255,32 @@ def test_get_history(monkeypatch, capsys):
         {"role": "user", "content": "Hello, world!", "model": "mock-model"},
         {"role": "assistant", "content": "I am a mock model!", "model": "mock-model"},
     ]
+
+
+def test_authenticate_method(monkeypatch):
+    # Mock the login.Auth class
+    mock_auth = MagicMock()
+    mock_auth.api_key = "test-api-key"
+    mock_auth.user_id = "test-user-id"
+
+    def mock_auth_constructor():
+        return mock_auth
+
+    monkeypatch.setattr("litai.client.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")
+
+    # Verify that the authentication was called
+    mock_auth.authenticate.assert_called_once()
+
+    # Verify that environment variables were set
+    assert os.getenv("LIGHTNING_API_KEY") == "test-api-key"
+    assert os.getenv("LIGHTNING_USER_ID") == "test-user-id"
+
+    # Test case 2: Neither api_key nor user_id provided
+    mock_auth.reset_mock()
+    LLM(model="openai/gpt-4")
+
+    # Verify that authentication was not called
+    mock_auth.authenticate.assert_not_called()