feat: LLMW V2 client changes

.github/workflows/ci_test.yaml

@@ -35,8 +35,8 @@ jobs:
     - name: Create test env
       shell: bash
       run: |
-        cp tests/sample.env tests/.env
-        sed -i "s|LLMWHISPERER_API_KEY=|LLMWHISPERER_API_KEY=${{ secrets.LLMWHISPERER_API_KEY }}|" tests/.env
+        cp sample.env .env
+        sed -i "s|LLMWHISPERER_API_KEY=|LLMWHISPERER_API_KEY=${{ secrets.LLMWHISPERER_API_KEY }}|" .env
 
     - name: Run tox
       id: tox

.pre-commit-config.yaml

@@ -17,6 +17,7 @@ repos:
         exclude_types:
           - "markdown"
       - id: end-of-file-fixer
+        exclude: "tests/test_data/.*"
       - id: check-yaml
         args: [--unsafe]
       - id: check-added-large-files
@@ -65,9 +66,7 @@ repos:
         args: [--max-line-length=120]
         exclude: |
           (?x)^(
-              .*migrations/.*\.py|
-              unstract-core/tests/.*|
-              pkgs/unstract-flags/src/unstract/flags/evaluation_.*\.py|
+              tests/test_data/.*|
           )$
   - repo: https://github.com/pycqa/isort
     rev: 5.13.2

pyproject.toml

@@ -69,7 +69,7 @@ includes = ["src"]
 package-dir = "src"
 
 [tool.pytest.ini_options]
-env_files = ["tests/.env"]
+env_files = [".env"]
 addopts = "-s"
 log_level = "INFO"
 log_cli = true

tests/sample.env → sample.env

@@ -1,3 +1,4 @@
 LLMWHISPERER_BASE_URL=https://llmwhisperer-api.unstract.com/v1
+LLMWHISPERER_BASE_URL_V2=https://llmwhisperer-api.us-central.unstract.com/api/v2
 LLMWHISPERER_LOG_LEVEL=DEBUG
 LLMWHISPERER_API_KEY=

src/unstract/llmwhisperer/__init__.py

@@ -1,8 +1,9 @@
 __version__ = "0.22.0"
 
 from .client import LLMWhispererClient  # noqa: F401
+from .client_v2 import LLMWhispererClientV2  # noqa: F401
 
 
-def get_sdk_version():
+def get_llmw_py_client_version():
     """Returns the SDK version."""
     return __version__

src/unstract/llmwhisperer/client.py

@@ -58,9 +58,7 @@ class LLMWhispererClient:
     client's activities and errors.
     """
 
-    formatter = logging.Formatter(
-        "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
-    )
+    formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
     logger = logging.getLogger(__name__)
     log_stream_handler = logging.StreamHandler()
     log_stream_handler.setFormatter(formatter)
@@ -117,9 +115,7 @@ def __init__(
             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.logger.debug("api_key set to %s", LLMWhispererUtils.redact_key(self.api_key))
 
         self.api_timeout = api_timeout
 
@@ -169,7 +165,7 @@ def whisper(
         ocr_provider: str = "advanced",
         line_splitter_tolerance: float = 0.4,
         horizontal_stretch_factor: float = 1.0,
-        encoding: str = "utf-8"
+        encoding: str = "utf-8",
     ) -> dict:
         """
         Sends a request to the LLMWhisperer API to process a document.
@@ -240,12 +236,10 @@ def whisper(
         should_stream = False
         if url == "":
             if stream is not None:
-
                 should_stream = True
 
                 def generate():
-                    for chunk in stream:
-                        yield chunk
+                    yield from stream
 
                 req = requests.Request(
                     "POST",
@@ -269,7 +263,7 @@ def generate():
             req = requests.Request("POST", api_url, params=params, headers=self.headers)
         prepared = req.prepare()
         s = requests.Session()
-        response = s.send(prepared, timeout=self.api_timeout, stream=should_stream)
+        response = s.send(prepared, timeout=timeout, stream=should_stream)
         response.encoding = encoding
         if response.status_code != 200 and response.status_code != 202:
             message = json.loads(response.text)