Merge pull request #40 from predictionguard/jacob/audio-doc-exts

Adding Extensions to Doc and Audio services, along with new Doc parameters

Author: jmansdorfer

predictionguard/src/audio.py

@@ -1,5 +1,3 @@
-import json
-
 import requests
 from typing import Any, Dict, Optional
 
@@ -46,6 +44,10 @@ def create(
         language: Optional[str] = "auto",
         temperature: Optional[float] = 0.0,
         prompt: Optional[str] = "",
+        toxicity: Optional[bool] = False,
+        pii: Optional[str] = "",
+        replace_method: Optional[str] = "",
+        injection: Optional[bool] = False,
     ) -> Dict[str, Any]:
         """
         Creates a audio transcription request to the Prediction Guard /audio/transcriptions API
@@ -55,25 +57,41 @@ def create(
         :param language: The language of the audio file
         :param temperature: The temperature parameter for model transcription
         :param prompt: A prompt to assist in transcription styling
+        :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.
+        :param injection: Whether to check for prompt injection
         :result: A dictionary containing the transcribed text.
         """
 
         # Create a list of tuples, each containing all the parameters for
         # a call to _transcribe_audio
-        args = (model, file, language, temperature, prompt)
+        args = (
+            model, file, language, temperature,
+            prompt, toxicity, pii, replace_method,
+            injection
+        )
 
         # Run _transcribe_audio
         choices = self._transcribe_audio(*args)
         return choices
 
-    def _transcribe_audio(self, model, file, language, temperature, prompt):
+    def _transcribe_audio(
+            self, model, file,
+            language, temperature, prompt,
+            toxicity, pii, replace_method, injection
+    ):
         """
         Function to transcribe an audio file.
         """
 
         headers = {
             "Authorization": "Bearer " + self.api_key,
             "User-Agent": "Prediction Guard Python Client: " + __version__,
+            "Toxicity": str(toxicity),
+            "Pii": pii,
+            "Replace-Method": replace_method,
+            "Injection": str(injection)
         }
 
         with open(file, "rb") as audio_file:

predictionguard/src/documents.py

@@ -1,6 +1,3 @@
-import json
-from pyexpat import model
-
 import requests
 from typing import Any, Dict, Optional
 
@@ -10,7 +7,7 @@
 class Documents:
     """Documents allows you to extract text from various document file types.
 
-    Usage::
+    Usage:
 
         from predictionguard import PredictionGuard
 
@@ -39,34 +36,71 @@ def __init__(self, api_key, url):
 
     def create(
         self,
-        file: str
+        file: str,
+        embed_images: Optional[bool] = False,
+        output_format: Optional[str] = None,
+        chunk_document: Optional[bool] = False,
+        chunk_size: Optional[int] = None,
+        toxicity: Optional[bool] = False,
+        pii: Optional[str] = "",
+        replace_method: Optional[str] = "",
+        injection: Optional[bool] = False,
     ) -> Dict[str, Any]:
         """
         Creates a documents request to the Prediction Guard /documents/extract API
 
         :param file: Document to be parsed
+        :param embed_images: Whether to embed images into documents
+        :param output_format: Output format
+        :param chunk_document: Whether to chunk documents into chunks
+        :param chunk_size: Chunk size
+        :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.
+        :param injection: Whether to check for prompt injection
         :result: A dictionary containing the title, content, and length of the document.
         """
 
         # Run _extract_documents
-        choices = self._extract_documents(file)
+        choices = self._extract_documents(
+            file, embed_images, output_format,
+            chunk_document, chunk_size, toxicity,
+            pii, replace_method, injection
+        )
         return choices
 
-    def _extract_documents(self, file):
+    def _extract_documents(
+            self, file, embed_images,
+            output_format, chunk_document,
+            chunk_size, toxicity, pii,
+            replace_method, injection
+    ):
         """
         Function to extract a document.
         """
 
         headers = {
             "Authorization": "Bearer " + self.api_key,
             "User-Agent": "Prediction Guard Python Client: " + __version__,
+            "Toxicity": str(toxicity),
+            "Pii": pii,
+            "Replace-Method": replace_method,
+            "Injection": str(injection)
+        }
+
+        data = {
+            "embedImages": embed_images,
+            "outputFormat": output_format,
+            "chunkDocument": chunk_document,
+            "chunkSize": chunk_size,
         }
 
         with open(file, "rb") as doc_file:
             files = {"file": (file, doc_file)}
 
             response = requests.request(
-                "POST", self.url + "/documents/extract", headers=headers, files=files
+                "POST", self.url + "/documents/extract",
+                headers=headers, files=files, data=data
             )
 
         # If the request was successful, print the proxies.

predictionguard/version.py

@@ -1,2 +1,2 @@
 # Setting the package version
-__version__ = "2.8.1"
+__version__ = "2.8.2"