adding new models function and updating old functions

Author: jmansdorfer

predictionguard/src/chat.py

@@ -272,14 +272,24 @@ def stream_generator(url, headers, payload, stream):
         else:
             return return_dict(self.url, headers, payload)
 
-    def list_models(self) -> List[str]:
+    def list_models(self, type: Optional[str, None]) -> List[str]:
         # Get the list of current models.
         headers = {
                 "Content-Type": "application/json",
                 "Authorization": "Bearer " + self.api_key,
                 "User-Agent": "Prediction Guard Python Client: " + __version__
                 }
 
-        response = requests.request("GET", self.url + "/chat/completions", headers=headers)
+        if type is None:
+            models_path = "/models/completion-chat"
+        else:
+            if type != "completion-chat" and type != "vision":
+                raise ValueError(
+                    "Please enter a valid models type (completion-chat or vision)."
+                )
+            else:
+                model_path = "/models/" + type
+
+        response = requests.request("GET", self.url + "/models/completion-chat", headers=headers)
 
         return list(response.json())

predictionguard/src/completions.py

@@ -110,6 +110,10 @@ def list_models(self) -> List[str]:
             "User-Agent": "Prediction Guard Python Client: " + __version__,
         }
 
-        response = requests.request("GET", self.url + "/completions", headers=headers)
+        response = requests.request("GET", self.url + "/models/completion", headers=headers)
 
-        return list(response.json())
+        response_list = []
+        for model in response.json()["data"]:
+            response_list.append(model)
+
+        return response_list

predictionguard/src/embeddings.py

@@ -4,7 +4,7 @@
 import base64
 
 import requests
-from typing import Any, Dict, List, Union
+from typing import Any, Dict, List, Union, Optional
 import urllib.request
 import urllib.parse
 import uuid
@@ -174,14 +174,29 @@ def _generate_embeddings(self, model, input, truncate, truncation_direction):
                 pass
             raise ValueError("Could not generate embeddings. " + err)
 
-    def list_models(self) -> List[str]:
+    def list_models(self, type: Optional[str, None] = None) -> List[str]:
         # Get the list of current models.
         headers = {
             "Content-Type": "application/json",
             "Authorization": "Bearer " + self.api_key,
             "User-Agent": "Prediction Guard Python Client: " + __version__,
         }
 
-        response = requests.request("GET", self.url + "/embeddings", headers=headers)
+        if type is None:
+            models_path = "/models/text-embeddings"
+        else:
+            if type != "text-embeddings" and type != "image-embeddings":
+                raise ValueError(
+                    "Please enter a valid models type "
+                    "(text-embeddings or image-embeddings)."
+                )
+            else:
+                models_path = "/models/" + type
+
+        response = requests.request("GET", self.url + models_path, headers=headers)
+
+        response_list = []
+        for model in response.json()["data"]:
+            response_list.append(model)
 
-        return list(response.json())
+        return response_list

predictionguard/src/models.py

@@ -0,0 +1,89 @@
+import requests
+from typing import Any, Dict, Optional
+
+from ..version import __version__
+
+
+class Models:
+    """Models lists all the models available in the Prediction Guard Platform.
+
+    Usage::
+
+        import os
+        import json
+
+        from predictionguard import PredictionGuard
+
+        # Set your Prediction Guard token as an environmental variable.
+        os.environ["PREDICTIONGUARD_API_KEY"] = "<api key>"
+
+        client = PredictionGuard()
+
+        response = client.models.list()
+
+        print(json.dumps(response, sort_keys=True, indent=4, separators=(",", ": ")))
+    """
+
+    def __init__(self, api_key, url):
+        self.api_key = api_key
+        self.url = url
+
+    def list(self, endpoint: Optional[str, None] = None) -> Dict[str, Any]:
+        """
+        Creates a models list request in the Prediction Guard REST API.
+
+        :param endpoint: The endpoint of models to list.
+        :return: A dictionary containing the metadata of all the models.
+        """
+
+        # Run _check_injection
+        choices = self._list_models(endpoint)
+        return choices
+
+    def _list_models(self, endpoint):
+        """
+        Function to list available models.
+        """
+
+        endpoints = [
+            "completion-chat", "completion", "vision",
+            "text-embeddings", "image-embeddings", "tokenize"
+        ]
+
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": "Bearer " + self.api_key,
+            "User-Agent": "Prediction Guard Python Client: " + __version__,
+        }
+
+        models_path = "/models"
+        if endpoint is not None:
+            if endpoint not in endpoints:
+                raise ValueError(
+                    "If specifying an endpoint, please use on of the following: "
+                    + ", ".join(endpoints)
+                )
+            else:
+                models_path += "/" + endpoint
+
+        response = requests.request(
+            "GET", self.url + models_path, headers=headers
+        )
+
+        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 check for injection. " + err)

predictionguard/src/tokenize.py

@@ -89,3 +89,19 @@ def _create_tokens(self, model, input):
             except Exception:
                 pass
             raise ValueError("Could not generate tokens. " + err)
+
+    def list_models(self):
+        # Get the list of current models.
+        headers = {
+                "Content-Type": "application/json",
+                "Authorization": "Bearer " + self.api_key,
+                "User-Agent": "Prediction Guard Python Client: " + __version__
+                }
+
+        response = requests.request("GET", self.url + "/models/tokenize", headers=headers)
+
+        response_list = []
+        for model in response.json()["data"]:
+            response_list.append(model)
+
+        return response_list

tests/test_chat.py

@@ -193,3 +193,15 @@ def test_chat_completions_list_models():
     response = test_client.chat.completions.list_models()
 
     assert len(response) > 0
+    assert type(response[0]) == str
+
+
+def test_chat_completions_list_models_fail():
+    test_client = PredictionGuard()
+
+    models_error = "Please enter a valid models type (completion-chat or vision)."
+
+    with pytest.raises(ValueError, match=models_error):
+        test_client.chat.completions.list_models(
+            type="fail"
+        )

tests/test_completions.py

@@ -32,3 +32,4 @@ def test_completions_list_models():
     response = test_client.completions.list_models()
 
     assert len(response) > 0
+    assert type(response[0]) == str

tests/test_embeddings.py

@@ -1,5 +1,4 @@
 import os
-import re
 import base64
 
 import pytest
@@ -215,3 +214,15 @@ def test_embeddings_list_models():
     response = test_client.embeddings.list_models()
 
     assert len(response) > 0
+    assert type(response[0]) is str
+
+
+def test_embeddings_list_models_fail():
+    test_client = PredictionGuard()
+
+    models_error = ""
+
+    with pytest.raises(ValueError, match=models_error):
+        test_client.embeddings.list_models(
+            type="fail"
+        )

tests/test_models.py

@@ -0,0 +1,90 @@
+from jedi.plugins import pytest
+from uaclient.api.u.pro.security.fix.cve.plan.v1 import endpoint
+
+from predictionguard import PredictionGuard
+
+
+def test_models_list():
+    test_client = PredictionGuard()
+
+    response = test_client.models.list()
+
+    assert len(response["data"]) > 0
+    assert type(response["data"][0]["id"]) is str
+
+
+def test_models_list_completion_chat():
+    test_client = PredictionGuard()
+
+    response = test_client.models.list(
+        endpoint="completion-chat"
+    )
+
+    assert len(response["data"]) > 0
+    assert type(response["data"][0]["id"]) is str
+
+
+def test_models_list_completion():
+    test_client = PredictionGuard()
+
+    response = test_client.models.list(
+        endpoint="completion"
+    )
+
+    assert len(response["data"]) > 0
+    assert type(response["data"][0]["id"]) is str
+
+
+def test_models_list_vision():
+    test_client = PredictionGuard()
+
+    response = test_client.models.list(
+        endpoint="vision"
+    )
+
+    assert len(response["data"]) > 0
+    assert type(response["data"][0]["id"]) is str
+
+
+def test_models_list_text_embeddings():
+    test_client = PredictionGuard()
+
+    response = test_client.models.list(
+        endpoint="text-embeddings"
+    )
+
+    assert len(response["data"]) > 0
+    assert type(response["data"][0]["id"]) is str
+
+
+def test_models_list_image_embeddings():
+    test_client = PredictionGuard()
+
+    response = test_client.models.list(
+        endpoint="image-embeddings"
+    )
+
+    assert len(response["data"]) > 0
+    assert type(response["data"][0]["id"]) is str
+
+
+def test_models_list_tokenize():
+    test_client = PredictionGuard()
+
+    response = test_client.models.list(
+        endpoint="tokenize"
+    )
+
+    assert len(response["data"]) > 0
+    assert type(response["data"][0]["id"]) is str
+
+
+def test_models_list_fail():
+    test_client = PredictionGuard()
+
+    models_error = ""
+
+    with pytest.raises(ValueError, match=models_error):
+        test_client.models.list(
+            endpoint="fail"
+        )

tests/test_tokenize.py

@@ -14,3 +14,11 @@ def test_tokenize_create():
     assert len(response) > 0
     assert type(response["tokens"][0]["id"]) is int
 
+
+def test_tokenize_list():
+    test_client = PredictionGuard()
+
+    response = test_client.tokenize.list_models()
+
+    assert len(response) > 0
+    assert type(response[0]) == str