From d2a63b76dae6f0e5e21b8f1f6639218c512c7f1d Mon Sep 17 00:00:00 2001 From: jmansdorfer Date: Fri, 1 Nov 2024 17:01:40 -0400 Subject: [PATCH 1/2] adding new models function and updating old functions --- predictionguard/src/chat.py | 14 ++++- predictionguard/src/completions.py | 8 ++- predictionguard/src/embeddings.py | 23 ++++++-- predictionguard/src/models.py | 89 +++++++++++++++++++++++++++++ predictionguard/src/tokenize.py | 16 ++++++ tests/test_chat.py | 12 ++++ tests/test_completions.py | 1 + tests/test_embeddings.py | 13 ++++- tests/test_models.py | 90 ++++++++++++++++++++++++++++++ tests/test_tokenize.py | 8 +++ 10 files changed, 265 insertions(+), 9 deletions(-) create mode 100644 predictionguard/src/models.py create mode 100644 tests/test_models.py diff --git a/predictionguard/src/chat.py b/predictionguard/src/chat.py index 289fd8d..c2dc789 100644 --- a/predictionguard/src/chat.py +++ b/predictionguard/src/chat.py @@ -272,7 +272,7 @@ 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", @@ -280,6 +280,16 @@ def list_models(self) -> List[str]: "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()) \ No newline at end of file diff --git a/predictionguard/src/completions.py b/predictionguard/src/completions.py index 1237991..9409504 100644 --- a/predictionguard/src/completions.py +++ b/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 diff --git a/predictionguard/src/embeddings.py b/predictionguard/src/embeddings.py index 3268fd1..ae201c4 100644 --- a/predictionguard/src/embeddings.py +++ b/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,7 +174,7 @@ 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", @@ -182,6 +182,21 @@ def list_models(self) -> List[str]: "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 diff --git a/predictionguard/src/models.py b/predictionguard/src/models.py new file mode 100644 index 0000000..8c78c86 --- /dev/null +++ b/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"] = "" + + 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) diff --git a/predictionguard/src/tokenize.py b/predictionguard/src/tokenize.py index 8bacd25..09f4219 100644 --- a/predictionguard/src/tokenize.py +++ b/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 diff --git a/tests/test_chat.py b/tests/test_chat.py index 717516d..9328ca4 100644 --- a/tests/test_chat.py +++ b/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" + ) diff --git a/tests/test_completions.py b/tests/test_completions.py index 3029a18..475c59a 100644 --- a/tests/test_completions.py +++ b/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 diff --git a/tests/test_embeddings.py b/tests/test_embeddings.py index 73aa24b..f75ccba 100644 --- a/tests/test_embeddings.py +++ b/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" + ) diff --git a/tests/test_models.py b/tests/test_models.py new file mode 100644 index 0000000..f4fcb01 --- /dev/null +++ b/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" + ) \ No newline at end of file diff --git a/tests/test_tokenize.py b/tests/test_tokenize.py index 6276a84..e67ef96 100644 --- a/tests/test_tokenize.py +++ b/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 From 4736b43a111dbc892173b15cda3edfebf30c547a Mon Sep 17 00:00:00 2001 From: jmansdorfer Date: Mon, 4 Nov 2024 14:54:59 -0500 Subject: [PATCH 2/2] adding new models endpoint function and updating old ones --- predictionguard/client.py | 6 ++++- predictionguard/src/chat.py | 24 ++++++++++---------- predictionguard/src/completions.py | 3 +-- predictionguard/src/embeddings.py | 21 ++++++++--------- predictionguard/src/models.py | 24 ++++++++++---------- predictionguard/src/tokenize.py | 2 +- tests/test_chat.py | 13 +---------- tests/test_completions.py | 2 +- tests/test_embeddings.py | 15 +------------ tests/test_models.py | 36 +++++++++--------------------- tests/test_tokenize.py | 2 +- 11 files changed, 55 insertions(+), 93 deletions(-) diff --git a/predictionguard/client.py b/predictionguard/client.py index 00b911b..a1e44f3 100644 --- a/predictionguard/client.py +++ b/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. diff --git a/predictionguard/src/chat.py b/predictionguard/src/chat.py index c2dc789..12532f8 100644 --- a/predictionguard/src/chat.py +++ b/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,7 +271,7 @@ 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", @@ -280,16 +279,17 @@ def list_models(self, type: Optional[str, None]) -> List[str]: "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()) \ No newline at end of file + return response_list \ No newline at end of file diff --git a/predictionguard/src/completions.py b/predictionguard/src/completions.py index 9409504..8a48f1d 100644 --- a/predictionguard/src/completions.py +++ b/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 diff --git a/predictionguard/src/embeddings.py b/predictionguard/src/embeddings.py index ae201c4..cdf1419 100644 --- a/predictionguard/src/embeddings.py +++ b/predictionguard/src/embeddings.py @@ -174,7 +174,7 @@ 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", @@ -182,21 +182,18 @@ def list_models(self, type: Optional[str, None] = None) -> List[str]: "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 diff --git a/predictionguard/src/models.py b/predictionguard/src/models.py index 8c78c86..c00a060 100644 --- a/predictionguard/src/models.py +++ b/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 diff --git a/predictionguard/src/tokenize.py b/predictionguard/src/tokenize.py index 09f4219..919422c 100644 --- a/predictionguard/src/tokenize.py +++ b/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 diff --git a/tests/test_chat.py b/tests/test_chat.py index 9328ca4..7c45a71 100644 --- a/tests/test_chat.py +++ b/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 diff --git a/tests/test_completions.py b/tests/test_completions.py index 475c59a..778583b 100644 --- a/tests/test_completions.py +++ b/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 diff --git a/tests/test_embeddings.py b/tests/test_embeddings.py index f75ccba..68edec0 100644 --- a/tests/test_embeddings.py +++ b/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 \ No newline at end of file diff --git a/tests/test_models.py b/tests/test_models.py index f4fcb01..80d9079 100644 --- a/tests/test_models.py +++ b/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" - ) \ No newline at end of file diff --git a/tests/test_tokenize.py b/tests/test_tokenize.py index e67ef96..39c04be 100644 --- a/tests/test_tokenize.py +++ b/tests/test_tokenize.py @@ -21,4 +21,4 @@ def test_tokenize_list(): response = test_client.tokenize.list_models() assert len(response) > 0 - assert type(response[0]) == str + assert type(response[0]) is str