adding new models endpoint function and updating old ones

Author: jmansdorfer

predictionguard/client.py

@@ -12,11 +12,12 @@
 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", "Tokenize",
-    "Translate", "Factuality", "Toxicity", "Pii", "Injection"
+    "Translate", "Factuality", "Toxicity", "Pii", "Injection", "Models"
 ]
 
 class PredictionGuard:
@@ -80,6 +81,9 @@ def __init__(
         self.tokenize: Tokenize = Tokenize(self.api_key, self.url)
         """Tokenize generates tokens for input text."""
 
+        self.models: Models = Models(self.api_key, self.url)
+        """Models lists all of the models available in the Prediction Guard API."""
+
     def _connect_client(self) -> None:
 
         # Prepare the proper headers.

predictionguard/src/chat.py

@@ -7,7 +7,6 @@
 from typing import Any, Dict, List, Optional, Union
 import urllib.request
 import urllib.parse
-from warnings import warn
 import uuid
 
 from ..version import __version__
@@ -272,24 +271,25 @@ def stream_generator(url, headers, payload, stream):
         else:
             return return_dict(self.url, headers, payload)
 
-    def list_models(self, type: Optional[str, None]) -> List[str]:
+    def list_models(self, capability: Optional[str] = "chat-completion") -> 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__
                 }
 
-        if type is None:
-            models_path = "/models/completion-chat"
+        if capability != "chat-completion" and capability != "chat-with-image":
+            raise ValueError(
+                "Please enter a valid model type (chat-completion or chat-with-image)."
+            )
         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
+            model_path = "/models/" + capability
+
+        response = requests.request("GET", self.url + model_path, headers=headers)
 
-        response = requests.request("GET", self.url + "/models/completion-chat", headers=headers)
+        response_list = []
+        for model in response.json()["data"]:
+            response_list.append(model["id"])
 
-        return list(response.json())
+        return response_list

predictionguard/src/completions.py

@@ -2,7 +2,6 @@
 
 import requests
 from typing import Any, Dict, List, Optional, Union
-from warnings import warn
 
 from ..version import __version__
 
@@ -114,6 +113,6 @@ def list_models(self) -> List[str]:
 
         response_list = []
         for model in response.json()["data"]:
-            response_list.append(model)
+            response_list.append(model["id"])
 
         return response_list

predictionguard/src/embeddings.py

@@ -174,29 +174,26 @@ def _generate_embeddings(self, model, input, truncate, truncation_direction):
                 pass
             raise ValueError("Could not generate embeddings. " + err)
 
-    def list_models(self, type: Optional[str, None] = None) -> List[str]:
+    def list_models(self, capability: Optional[str] = "embedding") -> 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__,
         }
 
-        if type is None:
-            models_path = "/models/text-embeddings"
+        if capability != "embedding" and capability != "embedding-with-image":
+            raise ValueError(
+                "Please enter a valid models type "
+                "(embedding or embedding-with-image)."
+            )
         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
+            model_path = "/models/" + capability
 
-        response = requests.request("GET", self.url + models_path, headers=headers)
+        response = requests.request("GET", self.url + model_path, headers=headers)
 
         response_list = []
         for model in response.json()["data"]:
-            response_list.append(model)
+            response_list.append(model["id"])
 
         return response_list

predictionguard/src/models.py

@@ -28,26 +28,26 @@ 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]:
+    def list(self, capability: Optional[str] = "") -> Dict[str, Any]:
         """
         Creates a models list request in the Prediction Guard REST API.
 
-        :param endpoint: The endpoint of models to list.
+        :param capability: The capability of models to list.
         :return: A dictionary containing the metadata of all the models.
         """
 
         # Run _check_injection
-        choices = self._list_models(endpoint)
+        choices = self._list_models(capability)
         return choices
 
-    def _list_models(self, endpoint):
+    def _list_models(self, capability):
         """
         Function to list available models.
         """
 
-        endpoints = [
-            "completion-chat", "completion", "vision",
-            "text-embeddings", "image-embeddings", "tokenize"
+        capabilities = [
+            "chat-completion", "chat-with-image", "completion",
+            "embedding", "embedding-with-image", "tokenize"
         ]
 
         headers = {
@@ -57,14 +57,14 @@ def _list_models(self, endpoint):
         }
 
         models_path = "/models"
-        if endpoint is not None:
-            if endpoint not in endpoints:
+        if capability != "":
+            if capability not in capabilities:
                 raise ValueError(
-                    "If specifying an endpoint, please use on of the following: "
-                    + ", ".join(endpoints)
+                    "If specifying a capability, please use one of the following: "
+                    + ", ".join(capabilities)
                 )
             else:
-                models_path += "/" + endpoint
+                models_path += "/" + capability
 
         response = requests.request(
             "GET", self.url + models_path, headers=headers

predictionguard/src/tokenize.py

@@ -102,6 +102,6 @@ def list_models(self):
 
         response_list = []
         for model in response.json()["data"]:
-            response_list.append(model)
+            response_list.append(model["id"])
 
         return response_list

tests/test_chat.py

@@ -193,15 +193,4 @@ 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"
-        )
+    assert type(response[0]) is str

tests/test_completions.py

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

tests/test_embeddings.py

@@ -1,8 +1,6 @@
 import os
 import base64
 
-import pytest
-
 from predictionguard import PredictionGuard
 
 
@@ -214,15 +212,4 @@ 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"
-        )
+    assert type(response[0]) is str

tests/test_models.py

@@ -1,6 +1,3 @@
-from jedi.plugins import pytest
-from uaclient.api.u.pro.security.fix.cve.plan.v1 import endpoint
-
 from predictionguard import PredictionGuard
 
 
@@ -13,55 +10,55 @@ def test_models_list():
     assert type(response["data"][0]["id"]) is str
 
 
-def test_models_list_completion_chat():
+def test_models_list_chat_completion():
     test_client = PredictionGuard()
 
     response = test_client.models.list(
-        endpoint="completion-chat"
+        capability="chat-completion"
     )
 
     assert len(response["data"]) > 0
     assert type(response["data"][0]["id"]) is str
 
 
-def test_models_list_completion():
+def test_models_list_chat_with_image():
     test_client = PredictionGuard()
 
     response = test_client.models.list(
-        endpoint="completion"
+        capability="chat-with-image"
     )
 
     assert len(response["data"]) > 0
     assert type(response["data"][0]["id"]) is str
 
 
-def test_models_list_vision():
+def test_models_list_completion():
     test_client = PredictionGuard()
 
     response = test_client.models.list(
-        endpoint="vision"
+        capability="completion"
     )
 
     assert len(response["data"]) > 0
     assert type(response["data"][0]["id"]) is str
 
 
-def test_models_list_text_embeddings():
+def test_models_list_embedding():
     test_client = PredictionGuard()
 
     response = test_client.models.list(
-        endpoint="text-embeddings"
+        capability="embedding"
     )
 
     assert len(response["data"]) > 0
     assert type(response["data"][0]["id"]) is str
 
 
-def test_models_list_image_embeddings():
+def test_models_list_embedding_with_image():
     test_client = PredictionGuard()
 
     response = test_client.models.list(
-        endpoint="image-embeddings"
+        capability="embedding-with-image"
     )
 
     assert len(response["data"]) > 0
@@ -72,19 +69,8 @@ def test_models_list_tokenize():
     test_client = PredictionGuard()
 
     response = test_client.models.list(
-        endpoint="tokenize"
+        capability="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"
-        )