adding enable_ocr param to doc service

.github/workflows/main.yml

@@ -20,12 +20,12 @@ jobs:
         run: pytest tests
         env:
           PREDICTIONGUARD_API_KEY: ${{ secrets.PREDICTIONGUARD_API_KEY }}
-          PREDICTIONGUARD_URL: ${{ secrets.PREDICTIONGUARD_URL }}
-          TEST_MODEL_NAME: ${{ secrets.TEST_MODEL_NAME }}
-          TEST_TEXT_EMBEDDINGS_MODEL: ${{ secrets.TEST_TEXT_EMBEDDINGS_MODEL }}
-          TEST_MULTIMODAL_EMBEDDINGS_MODEL: ${{ secrets.TEST_MULTIMODAL_EMBEDDINGS_MODEL }}
-          TEST_VISION_MODEL: ${{ secrets.TEST_VISION_MODEL }}
-          TEST_RERANK_MODEL: ${{ secrets.TEST_RERANK_MODEL }}
+          PREDICTIONGUARD_URL: ${{ vars.PREDICTIONGUARD_URL }}
+          TEST_CHAT_MODEL: ${{ vars.TEST_CHAT_MODEL }}
+          TEST_TEXT_EMBEDDINGS_MODEL: ${{ vars.TEST_TEXT_EMBEDDINGS_MODEL }}
+          TEST_MULTIMODAL_EMBEDDINGS_MODEL: ${{ vars.TEST_MULTIMODAL_EMBEDDINGS_MODEL }}
+          TEST_VISION_MODEL: ${{ vars.TEST_VISION_MODEL }}
+          TEST_RERANK_MODEL: ${{ vars.TEST_RERANK_MODEL }}
 
       - name: To PyPI using Flit
         uses: AsifArmanRahman/to-pypi-using-flit@v1

.github/workflows/pr.yml

@@ -24,9 +24,9 @@ jobs:
       run: pytest tests
       env:
         PREDICTIONGUARD_API_KEY: ${{ secrets.PREDICTIONGUARD_API_KEY}}
-        PREDICTIONGUARD_URL: ${{ secrets.PREDICTIONGUARD_URL}}
-        TEST_MODEL_NAME: ${{ secrets.TEST_MODEL_NAME }}
-        TEST_TEXT_EMBEDDINGS_MODEL: ${{ secrets.TEST_TEXT_EMBEDDINGS_MODEL }}
-        TEST_MULTIMODAL_EMBEDDINGS_MODEL: ${{ secrets.TEST_MULTIMODAL_EMBEDDINGS_MODEL }}
-        TEST_VISION_MODEL: ${{ secrets.TEST_VISION_MODEL }}
-        TEST_RERANK_MODEL: ${{ secrets.TEST_RERANK_MODEL }}
+        PREDICTIONGUARD_URL: ${{ vars.PREDICTIONGUARD_URL}}
+        TEST_CHAT_MODEL: ${{ vars.TEST_CHAT_MODEL }}
+        TEST_TEXT_EMBEDDINGS_MODEL: ${{ vars.TEST_TEXT_EMBEDDINGS_MODEL }}
+        TEST_MULTIMODAL_EMBEDDINGS_MODEL: ${{ vars.TEST_MULTIMODAL_EMBEDDINGS_MODEL }}
+        TEST_VISION_MODEL: ${{ vars.TEST_VISION_MODEL }}
+        TEST_RERANK_MODEL: ${{ vars.TEST_RERANK_MODEL }}

predictionguard/src/documents.py

@@ -41,6 +41,7 @@ def create(
         output_format: Optional[str] = None,
         chunk_document: Optional[bool] = False,
         chunk_size: Optional[int] = None,
+        enable_ocr: Optional[bool] = True,
         toxicity: Optional[bool] = False,
         pii: Optional[str] = "",
         replace_method: Optional[str] = "",
@@ -54,6 +55,7 @@ def create(
         :param output_format: Output format
         :param chunk_document: Whether to chunk documents into chunks
         :param chunk_size: Chunk size
+        :param enable_ocr: Whether to enable OCR
         :param toxicity: Whether to check for output toxicity
         :param pii: Whether to check for or replace pii
         :param replace_method: Replace method for any PII that is present.
@@ -64,16 +66,16 @@ def create(
         # Run _extract_documents
         choices = self._extract_documents(
             file, embed_images, output_format,
-            chunk_document, chunk_size, toxicity,
-            pii, replace_method, injection
+            chunk_document, chunk_size, enable_ocr,
+            toxicity, pii, replace_method, injection
         )
         return choices
 
     def _extract_documents(
             self, file, embed_images,
             output_format, chunk_document,
-            chunk_size, toxicity, pii,
-            replace_method, injection
+            chunk_size, enable_ocr, toxicity,
+            pii, replace_method, injection
     ):
         """
         Function to extract a document.
@@ -93,6 +95,7 @@ def _extract_documents(
             "outputFormat": output_format,
             "chunkDocument": chunk_document,
             "chunkSize": chunk_size,
+            "enableOCR": enable_ocr,
         }
 
         with open(file, "rb") as doc_file:

predictionguard/src/translate.py

@@ -7,116 +7,24 @@
 
 
 class Translate:
-    # UNCOMMENT WHEN DEPRECATED
-    # """No longer supported.
-    # """
-    #
-    # def __init__(self, api_key, url):
-    #     self.api_key = api_key
-    #     self.url = url
-    #
-    # def create(
-    #         self,
-    #         text: Optional[str],
-    #         source_lang: Optional[str],
-    #         target_lang: Optional[str],
-    #         use_third_party_engine: Optional[bool] = False
-    #     ) -> Dict[str, Any]:
-    #     """
-    #     No longer supported
-    #     """
-    #
-    #     raise ValueError(
-    #         "The translate functionality is no longer supported."
-    #     )
-    """Translate converts text from one language to another.
-
-    Usage::
-
-        from predictionguard import PredictionGuard
-
-        # Set your Prediction Guard token as an environmental variable.
-        os.environ["PREDICTIONGUARD_API_KEY"] = "<api key>"
-
-        client = PredictionGuard()
-
-        response = client.translate.create(
-            text="The sky is blue.",
-            source_lang="eng",
-            target_lang="fra",
-            use_third_party_engine=True
-        )
-
-        print(json.dumps(response, sort_keys=True, indent=4, separators=(",", ": ")))
+    """No longer supported.
     """
 
-    # REMOVE BELOW HERE FOR DEPRECATION
     def __init__(self, api_key, url):
         self.api_key = api_key
         self.url = url
 
     def create(
             self,
-            text: str,
-            source_lang: str,
-            target_lang: str,
+            text: Optional[str],
+            source_lang: Optional[str],
+            target_lang: Optional[str],
             use_third_party_engine: Optional[bool] = False
         ) -> Dict[str, Any]:
         """
-        Creates a translate request to the Prediction Guard /translate API.
-
-        :param text: The text to be translated.
-        :param source_lang: The language the text is currently in.
-        :param target_lang: The language the text will be translated to.
-        :param use_third_party_engine: A boolean for enabling translations with third party APIs.
-        :result: A dictionary containing the translate response.
-        """
-
-        # Create a list of tuples, each containing all the parameters for
-        # a call to _generate_translation
-        args = (text, source_lang, target_lang, use_third_party_engine)
-
-        # Run _generate_translation
-        choices = self._generate_translation(*args)
-        return choices
-
-    def _generate_translation(self, text, source_lang, target_lang, use_third_party_engine):
+        No longer supported
         """
-        Function to generate a translation response.
-        """
-
-        headers = {
-            "Content-Type": "application/json",
-            "Authorization": "Bearer " + self.api_key,
-            "User-Agent": "Prediction Guard Python Client: " + __version__,
-        }
-
-        payload_dict = {
-            "text": text,
-            "source_lang": source_lang,
-            "target_lang": target_lang,
-            "use_third_party_engine": use_third_party_engine
-        }
-        payload = json.dumps(payload_dict)
-        response = requests.request(
-            "POST", self.url + "/translate", headers=headers, data=payload
-        )
 
-        # If the request was successful, print the proxies.
-        if response.status_code == 200:
-            ret = response.json()
-            return ret
-        elif response.status_code == 429:
-            raise ValueError(
-                "Could not connect to Prediction Guard API. "
-                "Too many requests, rate limit or quota exceeded."
-            )
-        else:
-            # Check if there is a json body in the response. Read that in,
-            # print out the error field in the json body, and raise an exception.
-            err = ""
-            try:
-                err = response.json()["error"]
-            except Exception:
-                pass
-            raise ValueError("Could not make translation. " + err)
+        raise ValueError(
+            "The translate functionality is no longer supported."
+        )

tests/test_chat.py

@@ -10,7 +10,7 @@ def test_chat_completions_create():
     test_client = PredictionGuard()
 
     response = test_client.chat.completions.create(
-        model=os.environ["TEST_MODEL_NAME"],
+        model=os.environ["TEST_CHAT_MODEL"],
         messages=[
             {"role": "system", "content": "You are a helpful chatbot."},
             {"role": "user", "content": "Tell me a joke."},
@@ -24,7 +24,7 @@ def test_chat_completions_create_string():
     test_client = PredictionGuard()
 
     response = test_client.chat.completions.create(
-        model=os.environ["TEST_MODEL_NAME"],
+        model=os.environ["TEST_CHAT_MODEL"],
         messages="Tell me a joke"
     )
 
@@ -36,7 +36,7 @@ def test_chat_completions_create_stream():
 
     response_list = []
     for res in test_client.chat.completions.create(
-        model=os.environ["TEST_MODEL_NAME"],
+        model=os.environ["TEST_CHAT_MODEL"],
         messages=[
             {"role": "system", "content": "You are a helpful chatbot."},
             {"role": "user", "content": "Tell me a joke."},
@@ -58,7 +58,7 @@ def test_chat_completions_create_stream_output_fail():
     response_list = []
     with pytest.raises(ValueError, match=streaming_error):
         for res in test_client.chat.completions.create(
-            model=os.environ["TEST_MODEL_NAME"],
+            model=os.environ["TEST_CHAT_MODEL"],
             messages=[
                 {"role": "system", "content": "You are a helpful chatbot."},
                 {"role": "user", "content": "Tell me a joke."},
@@ -191,7 +191,7 @@ def test_chat_completions_create_tool_call():
     test_client = PredictionGuard()
 
     response = test_client.chat.completions.create(
-        model=os.environ["TEST_MODEL_NAME"],
+        model=os.environ["TEST_CHAT_MODEL"],
         messages=[
             {"role": "system", "content": "You are a helpful chatbot."},
             {"role": "user", "content": "Tell me a joke."},

tests/test_completions.py

@@ -9,7 +9,7 @@ def test_completions_create():
     test_client = PredictionGuard()
 
     response = test_client.completions.create(
-        model=os.environ["TEST_MODEL_NAME"], prompt="Tell me a joke"
+        model=os.environ["TEST_CHAT_MODEL"], prompt="Tell me a joke"
     )
 
     assert len(response["choices"][0]["text"]) > 0
@@ -19,7 +19,7 @@ def test_completions_create_batch():
     test_client = PredictionGuard()
 
     response = test_client.completions.create(
-        model=os.environ["TEST_MODEL_NAME"],
+        model=os.environ["TEST_CHAT_MODEL"],
         prompt=["Tell me a joke.", "Tell me a cool fact."],
     )
 
@@ -42,7 +42,7 @@ def test_completions_create_stream():
 
     response_list = []
     for res in test_client.completions.create(
-        model=os.environ["TEST_MODEL_NAME"],
+        model=os.environ["TEST_CHAT_MODEL"],
         prompt="Tell me a joke.",
         stream=True,
     ):
@@ -61,7 +61,7 @@ def test_completions_create_stream_output_fail():
     response_list = []
     with pytest.raises(ValueError, match=streaming_error):
         for res in test_client.completions.create(
-            model=os.environ["TEST_MODEL_NAME"],
+            model=os.environ["TEST_CHAT_MODEL"],
             prompt="Tell me a joke.",
             stream=True,
             output={"toxicity": True},

tests/test_tokenize.py

@@ -7,7 +7,7 @@ def test_tokenize_create():
     test_client = PredictionGuard()
 
     response = test_client.tokenize.create(
-        model=os.environ["TEST_MODEL_NAME"],
+        model=os.environ["TEST_CHAT_MODEL"],
         input="Tokenize this please."
     )