predictionguard 2.10.0__py3-none-any.whl

predictionguard/__init__.py

@@ -0,0 +1,6 @@
+"""Create controlled and compliant AI systems with PredictionGuard."""
+
+from .client import PredictionGuard as PredictionGuard
+from .version import __version__
+
+__version__ = __version__

predictionguard/client.py

@@ -0,0 +1,151 @@
+import os
+
+import requests
+from typing import Optional, Union
+
+from .src.audio import Audio
+from .src.chat import Chat
+from .src.completions import Completions
+from .src.detokenize import Detokenize
+from .src.documents import Documents
+from .src.embeddings import Embeddings
+from .src.rerank import Rerank
+from .src.tokenize import Tokenize
+from .src.translate import Translate
+from .src.factuality import Factuality
+from .src.toxicity import Toxicity
+from .src.pii import Pii
+from .src.injection import Injection
+from .src.models import Models
+from .version import __version__
+
+__all__ = [
+    "PredictionGuard", "Chat", "Completions", "Embeddings",
+    "Audio", "Documents", "Rerank", "Tokenize", "Translate",
+    "Detokenize", "Factuality", "Toxicity", "Pii", "Injection",
+    "Models"
+]
+
+class PredictionGuard:
+    """PredictionGuard provides access the Prediction Guard API."""
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        url: Optional[str] = None,
+        timeout: Optional[Union[int, float]] = None
+    ) -> None:
+        """
+        :param api_key: api_key represents PG api key.
+        :param url: url represents the transport and domain:port
+        :param timeout: request timeout in seconds.
+        """
+
+        # Get the access api_key.
+        if not api_key:
+            api_key = os.environ.get("PREDICTIONGUARD_API_KEY")
+
+        if not api_key:
+            raise ValueError(
+                "No api_key provided or in environment. "
+                "Please provide the api_key as "
+                "client = PredictionGuard(api_key=<your_api_key>) "
+                "or as PREDICTIONGUARD_API_KEY in your environment."
+            )
+        self.api_key = api_key
+
+        if not url:
+            url = os.environ.get("PREDICTIONGUARD_URL")
+        if not url:
+            url = "https://api.predictionguard.com"
+        self.url = url
+
+        if not timeout:
+            timeout = os.environ.get("TIMEOUT")
+        if not timeout:
+            timeout = None
+        if timeout:
+            try:
+                timeout = float(timeout)
+            except ValueError:
+                raise ValueError(
+                    "Timeout must be of type integer or float, not %s." % (type(timeout).__name__,)
+                )
+            except TypeError:
+                raise TypeError(
+                    "Timeout should be of type integer or float, not %s." % (type(timeout).__name__,)
+                )
+        self.timeout = timeout
+
+        # Connect to Prediction Guard and set the access api_key.
+        self._connect_client()
+
+        # Pass Prediction Guard class variables to inner classes
+        self.chat: Chat = Chat(self.api_key, self.url, self.timeout)
+        """Chat generates chat completions based on a conversation history"""
+
+        self.completions: Completions = Completions(self.api_key, self.url, self.timeout)
+        """Completions generates text completions based on the provided input"""
+
+        self.embeddings: Embeddings = Embeddings(self.api_key, self.url, self.timeout)
+        """Embedding generates chat completions based on a conversation history."""
+
+        self.audio: Audio = Audio(self.api_key, self.url, self.timeout)
+        """Audio allows for the transcription of audio files."""
+
+        self.documents: Documents = Documents(self.api_key, self.url, self.timeout)
+        """Documents allows you to extract text from various document file types."""
+
+        self.rerank: Rerank = Rerank(self.api_key, self.url, self.timeout)
+        """Rerank sorts text inputs by semantic relevance to a specified query."""
+
+        self.translate: Translate = Translate(self.api_key, self.url, self.timeout)
+        """Translate converts text from one language to another."""
+
+        self.factuality: Factuality = Factuality(self.api_key, self.url, self.timeout)
+        """Factuality checks the factuality of a given text compared to a reference."""
+
+        self.toxicity: Toxicity = Toxicity(self.api_key, self.url, self.timeout)
+        """Toxicity checks the toxicity of a given text."""
+
+        self.pii: Pii = Pii(self.api_key, self.url, self.timeout)
+        """Pii replaces personal information such as names, SSNs, and emails in a given text."""
+
+        self.injection: Injection = Injection(self.api_key, self.url, self.timeout)
+        """Injection detects potential prompt injection attacks in a given prompt."""
+
+        self.tokenize: Tokenize = Tokenize(self.api_key, self.url, self.timeout)
+        """Tokenize generates tokens for input text."""
+
+        self.detokenize: Detokenize = Detokenize(self.api_key, self.url, self.timeout)
+        """Detokenizes generates text for input tokens."""
+
+        self.models: Models = Models(self.api_key, self.url, self.timeout)
+        """Models lists all of the models available in the Prediction Guard API."""
+
+    def _connect_client(self) -> None:
+
+        # Prepare the proper headers.
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": "Bearer " + self.api_key,
+            "User-Agent": "Prediction Guard Python Client: " + __version__,
+        }
+
+        # Try listing models to make sure we can connect.
+        response = requests.request("GET", self.url + "/completions", headers=headers, timeout=self.timeout)
+
+        # If the connection was unsuccessful, raise an exception.
+        if response.status_code == 200:
+            pass
+        elif response.status_code == 401:
+            raise ValueError(
+                "Could not connect to Prediction Guard API with the given api_key. "
+                "Please check your access api_key and try again."
+            )
+        elif response.status_code == 404:
+            raise ValueError(
+                "Could not connect to Prediction Guard API with given url. "
+                "Please check url specified, if no url specified, "
+                "Please contact support."
+            )

predictionguard/src/audio.py

@@ -0,0 +1,194 @@
+from typing import Any, Dict, List, Optional
+
+import requests
+
+from ..version import __version__
+
+
+class Audio:
+    """
+    Audio generates a response based on audio data.
+
+    Usage::
+
+        import os
+        import json
+
+        from predictionguard import PredictionGuard
+
+        # Set your Prediction Guard token and url as an environmental variable.
+        os.environ["PREDICTIONGUARD_API_KEY"] = "<api key>"
+        os.environ["PREDICTIONGUARD_URL"] = "<url>"
+
+        # Or set your Prediction Guard token and url when initializing the PredictionGuard class.
+        client = PredictionGuard(
+            api_key="<api_key>",
+            url="<url>"
+        )
+
+        result = client.audio.transcriptions.create(
+            model="base",
+            file="sample_audio.wav"
+        )
+
+        print(json.dumps(
+            response,
+            sort_keys=True,
+            indent=4,
+            separators=(",", ": ")
+        ))
+    """
+
+    def __init__(self, api_key, url, timeout):
+        self.api_key = api_key
+        self.url = url
+        self.timeout = timeout
+
+        self.transcriptions: AudioTranscriptions = AudioTranscriptions(self.api_key, self.url, self.timeout)
+
+class AudioTranscriptions:
+    def __init__(self, api_key, url, timeout):
+        self.api_key = api_key
+        self.url = url
+        self.timeout = timeout
+
+    def create(
+        self,
+        model: str,
+        file: str,
+        language: Optional[str] = "auto",
+        temperature: Optional[float] = 0.0,
+        prompt: Optional[str] = "",
+        timestamp_granularities: Optional[List[str]] = None,
+        diarization: Optional[bool] = False,
+        response_format: Optional[str] = "json",
+        toxicity: Optional[bool] = False,
+        pii: Optional[str] = "",
+        replace_method: Optional[str] = "",
+        entity_list: Optional[List[str]] = None,
+        injection: Optional[bool] = False,
+    ) -> Dict[str, Any]:
+        """
+        Creates an audio transcription request to the Prediction Guard /audio/transcriptions API
+
+        :param model: The model to use
+        :param file: Audio file to be transcribed
+        :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 timestamp_granularities: The timestamp granularities to populate for this transcription
+        :param diarization: Whether to diarize the audio
+        :param response_format: The response format to use
+        :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 entity_list: List of entities to ignore in the PII check.
+        :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,
+            timestamp_granularities,
+            diarization,
+            response_format,
+            pii,
+            replace_method,
+            entity_list,
+            injection,
+            toxicity,
+        )
+
+        # Run _transcribe_audio
+        choices = self._transcribe_audio(*args)
+        return choices
+
+    def _transcribe_audio(
+            self,
+            model,
+            file,
+            language,
+            temperature,
+            prompt,
+            timestamp_granularities,
+            diarization,
+            response_format,
+            pii,
+            replace_method,
+            entity_list,
+            injection,
+            toxicity,
+    ):
+        """
+        Function to transcribe an audio file.
+        """
+
+        headers = {
+            "Authorization": "Bearer " + self.api_key,
+            "User-Agent": "Prediction Guard Python Client: " + __version__,
+            "Toxicity": str(toxicity),
+            "Pii": str(pii),
+            "Replace-Method": str(replace_method),
+            "Entity-List": str(entity_list),
+            "Injection": str(injection)
+        }
+
+        if timestamp_granularities:
+            if diarization and "segment" in timestamp_granularities:
+                raise ValueError(
+                    "Timestamp granularities cannot be set to "
+                    "`segments` when using diarization."
+                )
+
+            if response_format != "verbose_json":
+                raise ValueError(
+                    "Response format must be set to `verbose_json` "
+                    "when using timestamp granularities."
+                )
+
+        if diarization and response_format != "verbose_json":
+            raise ValueError(
+                "Response format must be set to `verbose_json` "
+                "when using diarization."
+            )
+
+        with open(file, "rb") as audio_file:
+            files = {"file": (file, audio_file, "audio/wav")}
+            data = {
+                "model": model,
+                "language": language,
+                "temperature": temperature,
+                "prompt": prompt,
+                "timestamp_granularities[]": timestamp_granularities,
+                "diarization": str(diarization).lower(),
+                "response_format": response_format,
+            }
+
+            response = requests.request(
+                "POST", self.url + "/audio/transcriptions", headers=headers, files=files, data=data, timeout=self.timeout
+            )
+
+        # 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 transcribe the audio file. " + err)