Merge pull request #2 from Zipstack/fix/assign-api-key-in-init

fix: Assign api_key on init call

.gitignore

@@ -159,5 +159,7 @@ cython_debug/
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 .idea/
 
+.vscode/
+
 .pdm-python
 .python-version

pyproject.toml

@@ -10,7 +10,7 @@ dependencies = [
     "requests~=2.32.3",
 ]
 readme = "README.md"
-urls = { Homepage = "https://llmwhisperer.unstract.com" }
+urls = { Homepage = "https://llmwhisperer.unstract.com", Source = "https://github.com/Zipstack/llm-whisperer-python-client" }
 license = {text = "AGPL v3"}
 authors = [
     {name = "Zipstack Inc", email = "[email protected]"},
@@ -29,6 +29,13 @@ classifiers = [
 ]
 
 [tool.pdm.dev-dependencies]
+test = [
+    "pytest>=8.2.2",
+    "pytest-mock>=3.14.0",
+    "pytest-dotenv>=0.5.2",
+    "pytest-cov>=5.0.0",
+    "pytest-md-report>=0.6.2",
+]
 lint = [
     "autopep8~=2.0.2",
     "black~=23.3.0",
@@ -41,6 +48,10 @@ lint = [
     "mypy~=1.10.0"
 ]
 
+[tool.pdm.version]
+source = "file"
+path = "src/llmwhisperer/__init__.py"
+
 [tool.isort]
 line_length = 120
 multi_line_output = 3
@@ -57,7 +68,13 @@ max-line-length = 120
 includes = ["src"]
 package-dir = "src"
 
-[tool.pdm.version]
-source = "file"
-path = "src/llmwhisperer/__init__.py"
+[tool.pytest.ini_options]
+env_files = ["tests/.env"]
+addopts = "-s"
+log_level = "INFO"
+log_cli = true
 
+[tool.pdm.scripts]
+test.cmd = "pytest"
+test.env_file = "tests/.env"
+test.help = "Runs pytests for LLM Whisperer client"

src/llmwhisperer/__init__.py

@@ -1,5 +1,4 @@
-__version__ = "0.11.0"
-
+__version__ = "0.1.1"
 
 def get_sdk_version():
     """Returns the SDK version."""

src/llmwhisperer/client.py

@@ -22,6 +22,8 @@
 
 import requests
 
+from llmwhisperer.utils import LLMWhispererUtils
+
 BASE_URL = "https://llmwhisperer-api.unstract.com/v1"
 
 
@@ -110,6 +112,9 @@ def __init__(
 
         if api_key == "":
             self.api_key = os.getenv("LLMWHISPERER_API_KEY", "")
+        else:
+            self.api_key = api_key
+        self.logger.debug("api_key set to %s", LLMWhispererUtils.redact_key(self.api_key))
 
         self.api_timeout = api_timeout
 

src/llmwhisperer/utils.py

@@ -0,0 +1,21 @@
+class LLMWhispererUtils:
+    @staticmethod
+    def redact_key(api_key: str, reveal_length=4) -> str:
+        """Hides sensitive information partially. Useful for logging keys.
+
+        Args:
+            api_key (str): API key to redact
+
+        Returns:
+            str: Redacted API key
+        """
+        if not isinstance(api_key, str):
+            raise ValueError("API key must be a string")
+
+        if reveal_length < 0:
+            raise ValueError("Reveal length must be a non-negative integer")
+
+        redacted_length = max(len(api_key) - reveal_length, 0)
+        revealed_part = api_key[:reveal_length]
+        redacted_part = "x" * redacted_length
+        return revealed_part + redacted_part

tests/sample.env

@@ -1,3 +1,3 @@
-export LLMWHISPERER_BASE_URL=https://llmwhisperer-api.unstract.com/v1
-export LLMWHISPERER_LOG_LEVEL=DEBUG
-export LLMWHISPERER_API_KEY=
+LLMWHISPERER_BASE_URL=https://llmwhisperer-api.unstract.com/v1
+LLMWHISPERER_LOG_LEVEL=DEBUG
+LLMWHISPERER_API_KEY=

tests/test_client.py

@@ -1,16 +1,36 @@
+import logging
 import unittest
 
+import pytest
+
 from llmwhisperer.client import LLMWhispererClient
 
+logger = logging.getLogger(__name__)
 
-class TestLLMWhispererClient(unittest.TestCase):
-    @unittest.skip("Skipping test_get_usage_info")
-    def test_get_usage_info(self):
-        client = LLMWhispererClient()
-        usage_info = client.get_usage_info()
-        print(usage_info)
-        self.assertIsInstance(usage_info, dict)
 
+@pytest.fixture
+def llm_whisperer_client():
+    # Create an instance of the client
+    client = LLMWhispererClient()
+    return client
+
+
+def test_get_usage_info(llm_whisperer_client):
+    usage_info = llm_whisperer_client.get_usage_info()
+    logger.info(usage_info)
+    assert isinstance(usage_info, dict), "usage_info should be a dictionary"
+    expected_keys = [
+        "current_page_count",
+        "daily_quota",
+        "monthly_quota",
+        "overage_page_count",
+        "subscription_plan",
+        "today_page_count",
+    ]
+    assert set(usage_info.keys()) == set(expected_keys), f"usage_info {usage_info} does not contain the expected keys"
+
+
+class TestLLMWhispererClient(unittest.TestCase):
     @unittest.skip("Skipping test_whisper")
     def test_whisper(self):
         client = LLMWhispererClient()
@@ -22,30 +42,31 @@ def test_whisper(self):
             timeout=200,
             store_metadata_for_highlighting=True,
         )
-        print(response)
+        logger.info(response)
         # self.assertIsInstance(response, dict)
 
     @unittest.skip("Skipping test_whisper_status")
     def test_whisper_status(self):
         client = LLMWhispererClient()
         response = client.whisper_status(whisper_hash="7cfa5cbb|5f1d285a7cf18d203de7af1a1abb0a3a")
-        print(response)
+        logger.info(response)
         self.assertIsInstance(response, dict)
 
     @unittest.skip("Skipping test_whisper_retrieve")
     def test_whisper_retrieve(self):
         client = LLMWhispererClient()
         response = client.whisper_retrieve(whisper_hash="7cfa5cbb|5f1d285a7cf18d203de7af1a1abb0a3a")
-        print(response)
+        logger.info(response)
         self.assertIsInstance(response, dict)
 
+    @unittest.skip("Skipping test_whisper_highlight_data")
     def test_whisper_highlight_data(self):
         client = LLMWhispererClient()
         response = client.highlight_data(
             whisper_hash="9924d865|5f1d285a7cf18d203de7af1a1abb0a3a",
             search_text="Indiranagar",
         )
-        print(response)
+        logger.info(response)
         self.assertIsInstance(response, dict)