Merge pull request #8 from Zipstack/streaming-upload

Added streaming upload

Author: jaseemjaskp

src/unstract/llmwhisperer/client.py

@@ -19,6 +19,7 @@
 import json
 import logging
 import os
+from typing import IO
 
 import requests
 
@@ -57,7 +58,9 @@ 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)
@@ -114,7 +117,9 @@ 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
 
@@ -150,6 +155,7 @@ def get_usage_info(self) -> dict:
     def whisper(
         self,
         file_path: str = "",
+        stream: IO[bytes] = None,
         url: str = "",
         processing_mode: str = "ocr",
         output_mode: str = "line-printer",
@@ -170,6 +176,7 @@ def whisper(
 
         Args:
             file_path (str, optional): The path to the file to be processed. Defaults to "".
+            stream (IO[bytes], optional): A stream of bytes to be processed. Defaults to None.
             url (str, optional): The URL of the file to be processed. Defaults to "".
             processing_mode (str, optional): The processing mode. Can be "ocr" or "text". Defaults to "ocr".
             output_mode (str, optional): The output mode. Can be "line-printer" or "text". Defaults to "line-printer".
@@ -212,11 +219,11 @@ def whisper(
         self.logger.debug("api_url: %s", api_url)
         self.logger.debug("params: %s", params)
 
-        if url == "" and file_path == "":
+        if url == "" and file_path == "" and stream is None:
             raise LLMWhispererClientException(
                 {
                     "status_code": -1,
-                    "message": "Either url or file_path must be provided",
+                    "message": "Either url, stream or file_path must be provided",
                 }
             )
 
@@ -228,21 +235,39 @@ def whisper(
                 }
             )
 
+        should_stream = False
         if url == "":
-            with open(file_path, "rb") as f:
-                data = f.read()
-            req = requests.Request(
-                "POST",
-                api_url,
-                params=params,
-                headers=self.headers,
-                data=data,
-            )
+            if stream is not None:
+
+                should_stream = True
+
+                def generate():
+                    for chunk in stream:
+                        yield chunk
+
+                req = requests.Request(
+                    "POST",
+                    api_url,
+                    params=params,
+                    headers=self.headers,
+                    data=generate(),
+                )
+
+            else:
+                with open(file_path, "rb") as f:
+                    data = f.read()
+                req = requests.Request(
+                    "POST",
+                    api_url,
+                    params=params,
+                    headers=self.headers,
+                    data=data,
+                )
         else:
             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)
+        response = s.send(prepared, timeout=self.api_timeout, stream=should_stream)
         if response.status_code != 200 and response.status_code != 202:
             message = json.loads(response.text)
             message["status_code"] = response.status_code

tests/client_test.py

@@ -4,6 +4,7 @@
 from pathlib import Path
 
 import pytest
+import requests
 
 from unstract.llmwhisperer import LLMWhispererClient
 
@@ -22,7 +23,9 @@ def test_get_usage_info(client):
         "subscription_plan",
         "today_page_count",
     ]
-    assert set(usage_info.keys()) == set(expected_keys), f"usage_info {usage_info} does not contain the expected keys"
+    assert set(usage_info.keys()) == set(
+        expected_keys
+    ), f"usage_info {usage_info} does not contain the expected keys"
 
 
 @pytest.mark.parametrize(
@@ -65,24 +68,45 @@ def test_whisper(self):
         #     url="https://storage.googleapis.com/pandora-static/samples/bill.jpg.pdf"
         # )
         response = client.whisper(
-            file_path="test_files/restaurant_invoice_photo.pdf",
+            file_path="test_data/restaurant_invoice_photo.pdf",
             timeout=200,
             store_metadata_for_highlighting=True,
         )
-        logger.info(response)
+        print(response)
+        # self.assertIsInstance(response, dict)
+
+    # @unittest.skip("Skipping test_whisper")
+    def test_whisper_stream(self):
+        client = LLMWhispererClient()
+        download_url = (
+            "https://storage.googleapis.com/pandora-static/samples/bill.jpg.pdf"
+        )
+        # Create a stream of download_url and pass it to whisper
+        response_download = requests.get(download_url, stream=True)
+        response_download.raise_for_status()
+        response = client.whisper(
+            stream=response_download.iter_content(chunk_size=1024),
+            timeout=200,
+            store_metadata_for_highlighting=True,
+        )
+        print(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")
+        response = client.whisper_status(
+            whisper_hash="7cfa5cbb|5f1d285a7cf18d203de7af1a1abb0a3a"
+        )
         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")
+        response = client.whisper_retrieve(
+            whisper_hash="7cfa5cbb|5f1d285a7cf18d203de7af1a1abb0a3a"
+        )
         logger.info(response)
         self.assertIsInstance(response, dict)