From 7c1120d7a4c5da9b3e7b7265a5c3784a0053db30 Mon Sep 17 00:00:00 2001 From: iamnotcj <69222710+iamnotcj@users.noreply.github.com> Date: Sat, 4 Jan 2025 12:04:06 -0600 Subject: [PATCH 1/9] WatsonX Generator Added watsonx generator --- docs/source/garak.generators.watsonx.rst | 7 ++ docs/source/generators.rst | 1 + garak/generators/watsonx.py | 108 +++++++++++++++++++++++ pyproject.toml | 1 + requirements.txt | 1 + 5 files changed, 118 insertions(+) create mode 100644 docs/source/garak.generators.watsonx.rst create mode 100644 garak/generators/watsonx.py diff --git a/docs/source/garak.generators.watsonx.rst b/docs/source/garak.generators.watsonx.rst new file mode 100644 index 000000000..90236dd72 --- /dev/null +++ b/docs/source/garak.generators.watsonx.rst @@ -0,0 +1,7 @@ +garak.generators.watsonx +======================= + +.. automodule:: garak.generators.watsonx + :members: + :undoc-members: + :show-inheritance: \ No newline at end of file diff --git a/docs/source/generators.rst b/docs/source/generators.rst index 0b7769b74..f74d91538 100644 --- a/docs/source/generators.rst +++ b/docs/source/generators.rst @@ -30,4 +30,5 @@ For a detailed oversight into how a generator operates, see :ref:`garak.generato garak.generators.rest garak.generators.rasa garak.generators.test + garak.generators.watsonx diff --git a/garak/generators/watsonx.py b/garak/generators/watsonx.py new file mode 100644 index 000000000..979bcaf90 --- /dev/null +++ b/garak/generators/watsonx.py @@ -0,0 +1,108 @@ +from garak import _config +from garak.generators.base import Generator +from typing import List, Union +import os +import importlib + + +class WatsonXGenerator(Generator): + """ + This is a generator for watsonx.ai. + + Make sure that you initialize the environment variables: 'WATSONX_TOKEN', 'WATSONX_URL', and 'WATSONX_PROJECTID'. + """ + ENV_VAR = "WATSONX_TOKEN" + URI_ENV_VAR = "WATSONX_URL" + PID_ENV_VAR = "WATSONX_PROJECTID" + DID_ENV_VAR = "WATSONX_DEPLOYID" + DEFAULT_PARAMS = Generator.DEFAULT_PARAMS | { + "role": "user", + "url": None, + "project_id": None, + "deployment_id": None, + "frequency_penalty": 0.5, + "logprobs": True, + "top_logprobs": 3, + "presence_penalty": 0.3, + "temperature": 0.7, + "max_tokens": 100, + "time_limit": 300000, + "top_p": 0.9, + "n": 1, + } + + generator_family_name = "watsonx" + + def __init__(self, name="", config_root=_config): + super().__init__(name, config_root=config_root) + + # Initialize and validate api_key + if self.api_key is not None: + os.environ[self.ENV_VAR] = self.api_key + + # Initialize and validate url. + self.url = os.getenv("WATSONX_URL", None) + if self.url is None: + raise ValueError( + f"The {self.URI_ENV_VAR} environment variable is required. Please enter the URL corresponding to the region of your provisioned service instance. \n" + ) + # Initialize and validate project_id. + self.project_id = os.getenv("WATSONX_PROJECTID", None) + if self.project_id is None: + raise ValueError( + f"The {self.PID_ENV_VAR} environment variable is required. Please enter the corresponding Project ID of the resource. \n" + ) + + # Import Foundation Models from ibm_watsonx_ai module. Import the Credentials function from the same module. + self.watsonx = importlib.import_module("ibm_watsonx_ai.foundation_models") + self.Credentials = getattr( + importlib.import_module("ibm_watsonx_ai"), "Credentials" + ) + def get_model(self): + # Call Credentials function with the url and api_key. + credentials = self.Credentials(url=self.url, api_key=self.api_key) + if self.name == "deployment/deployment": + self.deployment_id = os.getenv("WATSONX_DEPLOYID", None) + if self.deployment_id is None: + raise ValueError( + f"The {self.DID_ENV_VAR} environment variable is required. Please enter the corresponding Deployment ID of the resource. \n" + ) + + return self.watsonx.ModelInference( + deployment_id = self.deployment_id, + credentials=credentials, + project_id=self.project_id + ) + + else : + return self.watsonx.ModelInference( + model_id=self.name, + credentials=credentials, + project_id=self.project_id, + params=self.watsonx.schema.TextChatParameters( + frequency_penalty=self.frequency_penalty, + logprobs=self.logprobs, + top_logprobs=self.top_logprobs, + presence_penalty=self.presence_penalty, + temperature=self.temperature, + max_tokens=self.max_tokens, + time_limit=self.time_limit, + top_p=self.top_p, + n=self.n + ), + ) + + def _call_model(self, prompt: str, generations_this_call: int = 1) -> List[Union[str, None]]: + + # Get/Create Model + model = self.get_model() + + # Check if message is empty. If it is, append null byte. + if not prompt: + prompt = "\x00" + print("WARNING: Empty prompt was found. Null byte character appended to prevent API failure.") + + # Parse the output to only contain the output message from the model. Return a list containing that message. + return ["".join(model.generate(prompt=prompt)['results'][0]['generated_text'])] + +DEFAULT_CLASS = "WatsonXGenerator" \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index eccd8a931..41ab46e0a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -75,6 +75,7 @@ dependencies = [ "zalgolib>=0.2.2", "ecoji>=0.1.1", "deepl==1.17.0", + "ibm-watsonx-ai==1.1.25", "fschat>=0.2.36", "litellm>=1.41.21", "jsonpath-ng>=1.6.1", diff --git a/requirements.txt b/requirements.txt index 50de30fe5..58e6f43c1 100644 --- a/requirements.txt +++ b/requirements.txt @@ -28,6 +28,7 @@ deepl==1.17.0 fschat>=0.2.36 litellm>=1.41.21 jsonpath-ng>=1.6.1 +ibm-watsonx-ai==1.1.25 huggingface_hub>=0.21.0 python-magic-bin>=0.4.14; sys_platform == "win32" python-magic>=0.4.21; sys_platform != "win32" From 17703fc69e6888b5042a5ef644fd1ccebf3847f3 Mon Sep 17 00:00:00 2001 From: iamnotcj <69222710+iamnotcj@users.noreply.github.com> Date: Sat, 4 Jan 2025 12:07:20 -0600 Subject: [PATCH 2/9] Updated Format Black Format --- garak/generators/watsonx.py | 27 +++++++++++++++++---------- 1 file changed, 17 insertions(+), 10 deletions(-) diff --git a/garak/generators/watsonx.py b/garak/generators/watsonx.py index 979bcaf90..88eac9ec1 100644 --- a/garak/generators/watsonx.py +++ b/garak/generators/watsonx.py @@ -11,6 +11,7 @@ class WatsonXGenerator(Generator): Make sure that you initialize the environment variables: 'WATSONX_TOKEN', 'WATSONX_URL', and 'WATSONX_PROJECTID'. """ + ENV_VAR = "WATSONX_TOKEN" URI_ENV_VAR = "WATSONX_URL" PID_ENV_VAR = "WATSONX_PROJECTID" @@ -52,12 +53,13 @@ def __init__(self, name="", config_root=_config): raise ValueError( f"The {self.PID_ENV_VAR} environment variable is required. Please enter the corresponding Project ID of the resource. \n" ) - + # Import Foundation Models from ibm_watsonx_ai module. Import the Credentials function from the same module. self.watsonx = importlib.import_module("ibm_watsonx_ai.foundation_models") self.Credentials = getattr( importlib.import_module("ibm_watsonx_ai"), "Credentials" ) + def get_model(self): # Call Credentials function with the url and api_key. credentials = self.Credentials(url=self.url, api_key=self.api_key) @@ -67,14 +69,14 @@ def get_model(self): raise ValueError( f"The {self.DID_ENV_VAR} environment variable is required. Please enter the corresponding Deployment ID of the resource. \n" ) - + return self.watsonx.ModelInference( - deployment_id = self.deployment_id, + deployment_id=self.deployment_id, credentials=credentials, - project_id=self.project_id + project_id=self.project_id, ) - else : + else: return self.watsonx.ModelInference( model_id=self.name, credentials=credentials, @@ -88,11 +90,13 @@ def get_model(self): max_tokens=self.max_tokens, time_limit=self.time_limit, top_p=self.top_p, - n=self.n + n=self.n, ), ) - def _call_model(self, prompt: str, generations_this_call: int = 1) -> List[Union[str, None]]: + def _call_model( + self, prompt: str, generations_this_call: int = 1 + ) -> List[Union[str, None]]: # Get/Create Model model = self.get_model() @@ -100,9 +104,12 @@ def _call_model(self, prompt: str, generations_this_call: int = 1) -> List[Union # Check if message is empty. If it is, append null byte. if not prompt: prompt = "\x00" - print("WARNING: Empty prompt was found. Null byte character appended to prevent API failure.") + print( + "WARNING: Empty prompt was found. Null byte character appended to prevent API failure." + ) # Parse the output to only contain the output message from the model. Return a list containing that message. - return ["".join(model.generate(prompt=prompt)['results'][0]['generated_text'])] + return ["".join(model.generate(prompt=prompt)["results"][0]["generated_text"])] + -DEFAULT_CLASS = "WatsonXGenerator" \ No newline at end of file +DEFAULT_CLASS = "WatsonXGenerator" From 8149e90e3e3b36c9dc080d39f62b7a1d2522bbae Mon Sep 17 00:00:00 2001 From: iamnotcj <69222710+iamnotcj@users.noreply.github.com> Date: Sat, 4 Jan 2025 12:11:24 -0600 Subject: [PATCH 3/9] Minor update Added more documentation. --- garak/generators/watsonx.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/garak/generators/watsonx.py b/garak/generators/watsonx.py index 88eac9ec1..d78d65e42 100644 --- a/garak/generators/watsonx.py +++ b/garak/generators/watsonx.py @@ -9,7 +9,13 @@ class WatsonXGenerator(Generator): """ This is a generator for watsonx.ai. - Make sure that you initialize the environment variables: 'WATSONX_TOKEN', 'WATSONX_URL', and 'WATSONX_PROJECTID'. + Make sure that you initialize the environment variables: + 'WATSONX_TOKEN', + 'WATSONX_URL', + and 'WATSONX_PROJECTID'. + + To use a tuned model that is deployed, use 'deployment/deployment' for the -n flag and make sure + to also initialize the 'WATSONX_DEPLOYID' environment variable. """ ENV_VAR = "WATSONX_TOKEN" From c0b363138579bfa28655d0e1360cbfdac01df349 Mon Sep 17 00:00:00 2001 From: iamnotcj <69222710+iamnotcj@users.noreply.github.com> Date: Sat, 4 Jan 2025 15:39:23 -0600 Subject: [PATCH 4/9] Update Changed variable names and updated tests --- garak/generators/watsonx.py | 31 +++++++++++++++++------------ tests/generators/test_generators.py | 1 + 2 files changed, 19 insertions(+), 13 deletions(-) diff --git a/garak/generators/watsonx.py b/garak/generators/watsonx.py index d78d65e42..8edc49a05 100644 --- a/garak/generators/watsonx.py +++ b/garak/generators/watsonx.py @@ -23,8 +23,7 @@ class WatsonXGenerator(Generator): PID_ENV_VAR = "WATSONX_PROJECTID" DID_ENV_VAR = "WATSONX_DEPLOYID" DEFAULT_PARAMS = Generator.DEFAULT_PARAMS | { - "role": "user", - "url": None, + "uri": None, "project_id": None, "deployment_id": None, "frequency_penalty": 0.5, @@ -48,17 +47,23 @@ def __init__(self, name="", config_root=_config): os.environ[self.ENV_VAR] = self.api_key # Initialize and validate url. - self.url = os.getenv("WATSONX_URL", None) - if self.url is None: - raise ValueError( - f"The {self.URI_ENV_VAR} environment variable is required. Please enter the URL corresponding to the region of your provisioned service instance. \n" - ) + if self.uri is not None: + pass + else : + self.uri = os.getenv("WATSONX_URL", None) + if self.uri is None: + raise ValueError( + f"The {self.URI_ENV_VAR} environment variable is required. Please enter the URL corresponding to the region of your provisioned service instance. \n" + ) # Initialize and validate project_id. - self.project_id = os.getenv("WATSONX_PROJECTID", None) - if self.project_id is None: - raise ValueError( - f"The {self.PID_ENV_VAR} environment variable is required. Please enter the corresponding Project ID of the resource. \n" - ) + if self.project_id is not None: + pass + else : + self.project_id = os.getenv("WATSONX_PROJECTID", None) + if self.project_id is None: + raise ValueError( + f"The {self.PID_ENV_VAR} environment variable is required. Please enter the corresponding Project ID of the resource. \n" + ) # Import Foundation Models from ibm_watsonx_ai module. Import the Credentials function from the same module. self.watsonx = importlib.import_module("ibm_watsonx_ai.foundation_models") @@ -68,7 +73,7 @@ def __init__(self, name="", config_root=_config): def get_model(self): # Call Credentials function with the url and api_key. - credentials = self.Credentials(url=self.url, api_key=self.api_key) + credentials = self.Credentials(url=self.uri, api_key=self.api_key) if self.name == "deployment/deployment": self.deployment_id = os.getenv("WATSONX_DEPLOYID", None) if self.deployment_id is None: diff --git a/tests/generators/test_generators.py b/tests/generators/test_generators.py index 74c2a153c..6e55b5c66 100644 --- a/tests/generators/test_generators.py +++ b/tests/generators/test_generators.py @@ -211,6 +211,7 @@ def test_instantiate_generators(classname): "org_id": "fake", # required for NeMo "uri": "https://example.com", # required for rest "provider": "fake", # required for LiteLLM + "project_id": "fake", # required for watsonx } } } From 6c008b8eeb05f0ac5ebf393423ec415c623a51e7 Mon Sep 17 00:00:00 2001 From: iamnotcj <69222710+iamnotcj@users.noreply.github.com> Date: Fri, 17 Jan 2025 22:09:52 -0600 Subject: [PATCH 5/9] Update .gitignore --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index 809828e8a..abee52cbc 100644 --- a/.gitignore +++ b/.gitignore @@ -168,3 +168,4 @@ hitlog.*.jsonl garak_runs/ runs/ logs/ +.DS_Store From ecaa821b98fab2b17be3fd7230790d4c7423ac86 Mon Sep 17 00:00:00 2001 From: iamnotcj <69222710+iamnotcj@users.noreply.github.com> Date: Sat, 25 Jan 2025 01:26:15 -0600 Subject: [PATCH 6/9] 2nd Update Dropped the SDK for requests library, Added some tests --- garak/generators/watsonx.py | 163 ++++++++++++++++------------ pyproject.toml | 1 - requirements.txt | 1 - tests/generators/conftest.py | 6 + tests/generators/test_generators.py | 2 +- tests/generators/test_watsonx.py | 81 ++++++++++++++ tests/generators/watsonx.json | 29 +++++ 7 files changed, 210 insertions(+), 73 deletions(-) create mode 100644 tests/generators/test_watsonx.py create mode 100644 tests/generators/watsonx.json diff --git a/garak/generators/watsonx.py b/garak/generators/watsonx.py index 8edc49a05..49b767dde 100644 --- a/garak/generators/watsonx.py +++ b/garak/generators/watsonx.py @@ -2,20 +2,20 @@ from garak.generators.base import Generator from typing import List, Union import os -import importlib +import requests class WatsonXGenerator(Generator): """ This is a generator for watsonx.ai. - Make sure that you initialize the environment variables: - 'WATSONX_TOKEN', - 'WATSONX_URL', - and 'WATSONX_PROJECTID'. - - To use a tuned model that is deployed, use 'deployment/deployment' for the -n flag and make sure - to also initialize the 'WATSONX_DEPLOYID' environment variable. + Make sure that you initialize the environment variables: + 'WATSONX_TOKEN', + 'WATSONX_URL', + 'WATSONX_PROJECTID' OR 'WATSONX_DEPLOYID'. + + To use a model that is in the "project" stage initialize the WATSONX_PROJECTID variable with the Project ID of the model. + To use a tuned model that is deployed, simply initialize the WATSONX_DEPLOYID variable with the Deployment ID of the model. """ ENV_VAR = "WATSONX_TOKEN" @@ -24,93 +24,110 @@ class WatsonXGenerator(Generator): DID_ENV_VAR = "WATSONX_DEPLOYID" DEFAULT_PARAMS = Generator.DEFAULT_PARAMS | { "uri": None, - "project_id": None, - "deployment_id": None, - "frequency_penalty": 0.5, - "logprobs": True, - "top_logprobs": 3, - "presence_penalty": 0.3, - "temperature": 0.7, - "max_tokens": 100, - "time_limit": 300000, - "top_p": 0.9, - "n": 1, + "project_id": "", + "deployment_id": "", + "prompt_variable": "input", + "bearer_token": "", + "max_tokens": 900, } generator_family_name = "watsonx" def __init__(self, name="", config_root=_config): super().__init__(name, config_root=config_root) - # Initialize and validate api_key if self.api_key is not None: os.environ[self.ENV_VAR] = self.api_key - + + def _set_bearer_token(self, iam_url="https://iam.cloud.ibm.com/identity/token"): + header = { + "Content-Type": "application/x-www-form-urlencoded", + "Accept": "application/json", + } + body = ( + "grant_type=urn:ibm:params:oauth:grant-type:apikey&apikey=" + self.api_key + ) + response = requests.post(url=iam_url, headers=header, data=body) + self.bearer_token = "Bearer " + response.json()["access_token"] + + def _generate_with_project(self, payload): + # Generation via Project ID. + + url = self.uri + "/ml/v1/text/generation?version=2023-05-29" + + body = { + "input": payload, + "parameters": { + "decoding_method": "greedy", + "max_new_tokens": self.max_tokens, + "min_new_tokens": 0, + "repetition_penalty": 1, + }, + "model_id": self.name, + "project_id": self.project_id, + } + + headers = { + "Accept": "application/json", + "Content-Type": "application/json", + "Authorization": self.bearer_token, + } + + response = requests.post(url=url, headers=headers, json=body) + return response.json() + + def _generate_with_deployment(self, payload): + # Generation via Deployment ID. + url = ( + self.uri + + "/ml/v1/deployments/" + + self.deployment_id + + "/text/generation?version=2021-05-01" + ) + body = {"parameters": {"prompt_variables": {self.prompt_variable: payload}}} + headers = { + "Content-Type": "application/json", + "Accept": "application/json", + "Authorization": self.bearer_token, + } + response = requests.post(url=url, headers=headers, json=body) + return response.json() + + def _validate_env_var(self): # Initialize and validate url. if self.uri is not None: pass - else : + else: self.uri = os.getenv("WATSONX_URL", None) if self.uri is None: raise ValueError( f"The {self.URI_ENV_VAR} environment variable is required. Please enter the URL corresponding to the region of your provisioned service instance. \n" ) + # Initialize and validate project_id. - if self.project_id is not None: + if self.project_id: pass - else : - self.project_id = os.getenv("WATSONX_PROJECTID", None) - if self.project_id is None: - raise ValueError( - f"The {self.PID_ENV_VAR} environment variable is required. Please enter the corresponding Project ID of the resource. \n" - ) - - # Import Foundation Models from ibm_watsonx_ai module. Import the Credentials function from the same module. - self.watsonx = importlib.import_module("ibm_watsonx_ai.foundation_models") - self.Credentials = getattr( - importlib.import_module("ibm_watsonx_ai"), "Credentials" - ) - - def get_model(self): - # Call Credentials function with the url and api_key. - credentials = self.Credentials(url=self.uri, api_key=self.api_key) - if self.name == "deployment/deployment": - self.deployment_id = os.getenv("WATSONX_DEPLOYID", None) - if self.deployment_id is None: - raise ValueError( - f"The {self.DID_ENV_VAR} environment variable is required. Please enter the corresponding Deployment ID of the resource. \n" - ) - - return self.watsonx.ModelInference( - deployment_id=self.deployment_id, - credentials=credentials, - project_id=self.project_id, - ) + else: + self.project_id = os.getenv("WATSONX_PROJECTID", "") + # Initialize and validate deployment_id. + if self.deployment_id: + pass else: - return self.watsonx.ModelInference( - model_id=self.name, - credentials=credentials, - project_id=self.project_id, - params=self.watsonx.schema.TextChatParameters( - frequency_penalty=self.frequency_penalty, - logprobs=self.logprobs, - top_logprobs=self.top_logprobs, - presence_penalty=self.presence_penalty, - temperature=self.temperature, - max_tokens=self.max_tokens, - time_limit=self.time_limit, - top_p=self.top_p, - n=self.n, - ), + self.deployment_id = os.getenv("WATSONX_DEPLOYID", "") + + # Check to ensure at least ONE of project_id or deployment_id is populated. + if not self.project_id and not self.deployment_id: + raise ValueError( + f"Either {self.PID_ENV_VAR} or {self.DID_ENV_VAR} is required. Please supply either a Project ID or Deployment ID. \n" ) + return super()._validate_env_var() def _call_model( self, prompt: str, generations_this_call: int = 1 ) -> List[Union[str, None]]: - - # Get/Create Model - model = self.get_model() + if not self.bearer_token : + self._set_bearer_token() # Check if message is empty. If it is, append null byte. if not prompt: @@ -119,8 +136,14 @@ def _call_model( "WARNING: Empty prompt was found. Null byte character appended to prevent API failure." ) + output = "" + if self.deployment_id: + output = self._generate_with_deployment(prompt) + else: + output = self._generate_with_project(prompt) + # Parse the output to only contain the output message from the model. Return a list containing that message. - return ["".join(model.generate(prompt=prompt)["results"][0]["generated_text"])] + return ["".join(output["results"][0]["generated_text"])] DEFAULT_CLASS = "WatsonXGenerator" diff --git a/pyproject.toml b/pyproject.toml index 41ab46e0a..eccd8a931 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -75,7 +75,6 @@ dependencies = [ "zalgolib>=0.2.2", "ecoji>=0.1.1", "deepl==1.17.0", - "ibm-watsonx-ai==1.1.25", "fschat>=0.2.36", "litellm>=1.41.21", "jsonpath-ng>=1.6.1", diff --git a/requirements.txt b/requirements.txt index 58e6f43c1..50de30fe5 100644 --- a/requirements.txt +++ b/requirements.txt @@ -28,7 +28,6 @@ deepl==1.17.0 fschat>=0.2.36 litellm>=1.41.21 jsonpath-ng>=1.6.1 -ibm-watsonx-ai==1.1.25 huggingface_hub>=0.21.0 python-magic-bin>=0.4.14; sys_platform == "win32" python-magic>=0.4.21; sys_platform != "win32" diff --git a/tests/generators/conftest.py b/tests/generators/conftest.py index 9a760d80f..52d89c163 100644 --- a/tests/generators/conftest.py +++ b/tests/generators/conftest.py @@ -18,3 +18,9 @@ def hf_endpoint_mocks(): """Mock responses for Huggingface InferenceAPI based endpoints""" with open(pathlib.Path(__file__).parents[0] / "hf_inference.json") as mock_openai: return json.load(mock_openai) + +@pytest.fixture +def watsonx_compat_mocks(): + """Mock responses for watsonx.ai based endpoints""" + with open(pathlib.Path(__file__).parents[0] / "watsonx.json") as mock_watsonx: + return json.load(mock_watsonx) diff --git a/tests/generators/test_generators.py b/tests/generators/test_generators.py index 6e55b5c66..d8573d153 100644 --- a/tests/generators/test_generators.py +++ b/tests/generators/test_generators.py @@ -184,6 +184,7 @@ def test_generator_structure(classname): if classname not in [ "generators.azure.AzureOpenAIGenerator", # requires additional env variables tested in own test class + "generators.watsonx.WatsonXGenerator", # requires additional env variables tested in own test class "generators.function.Multiple", # requires mock local function not implemented here "generators.function.Single", # requires mock local function not implemented here "generators.ggml.GgmlGenerator", # validates files on disk tested in own test class @@ -211,7 +212,6 @@ def test_instantiate_generators(classname): "org_id": "fake", # required for NeMo "uri": "https://example.com", # required for rest "provider": "fake", # required for LiteLLM - "project_id": "fake", # required for watsonx } } } diff --git a/tests/generators/test_watsonx.py b/tests/generators/test_watsonx.py new file mode 100644 index 000000000..e9683d559 --- /dev/null +++ b/tests/generators/test_watsonx.py @@ -0,0 +1,81 @@ +from garak.generators.watsonx import WatsonXGenerator +import os +import pytest +import requests_mock + + +DEFAULT_DEPLOYMENT_NAME = "ibm/granite-3-8b-instruct" + + +@pytest.fixture +def set_fake_env(request) -> None: + stored_env = { + WatsonXGenerator.ENV_VAR: os.getenv(WatsonXGenerator.ENV_VAR, None), + WatsonXGenerator.PID_ENV_VAR: os.getenv(WatsonXGenerator.PID_ENV_VAR, None), + WatsonXGenerator.URI_ENV_VAR: os.getenv(WatsonXGenerator.URI_ENV_VAR, None), + WatsonXGenerator.DID_ENV_VAR: os.getenv(WatsonXGenerator.DID_ENV_VAR, None), + } + + def restore_env(): + for k, v in stored_env.items(): + if v is not None: + os.environ[k] = v + else: + del os.environ[k] + + os.environ[WatsonXGenerator.ENV_VAR] = "XXXXXXXXXXXXX" + os.environ[WatsonXGenerator.PID_ENV_VAR] = "XXXXXXXX-XXXX-XXXX-XXXX-XXXXXXXXXXXX" + os.environ[WatsonXGenerator.DID_ENV_VAR] = "XXXXXXXX-XXXX-XXXX-XXXX-XXXXXXXXXXXX" + os.environ[WatsonXGenerator.URI_ENV_VAR] = "https://garak.example.com/" + request.addfinalizer(restore_env) + + +@pytest.mark.usefixtures("set_fake_env") +def test_bearer_token(watsonx_compat_mocks): + with requests_mock.Mocker() as m: + mock_response = watsonx_compat_mocks["watsonx_bearer_token"] + + extended_request = "identity/token" + + m.post( + "https://garak.example.com/" + extended_request, json=mock_response["json"] + ) + + granite_llm = WatsonXGenerator(DEFAULT_DEPLOYMENT_NAME) + token = granite_llm._set_bearer_token(iam_url="https://garak.example.com/identity/token") + + assert granite_llm.bearer_token == ("Bearer " + mock_response["json"]["access_token"]) + + +@pytest.mark.usefixtures("set_fake_env") +def test_project(watsonx_compat_mocks): + with requests_mock.Mocker() as m: + mock_response = watsonx_compat_mocks["watsonx_generation"] + extended_request = "/ml/v1/text/generation?version=2023-05-29" + + m.post( + "https://garak.example.com/" + extended_request, json=mock_response["json"] + ) + + granite_llm = WatsonXGenerator(DEFAULT_DEPLOYMENT_NAME) + response = granite_llm._generate_with_project("What is this?") + + assert granite_llm.name == response["model_id"] + + +@pytest.mark.usefixtures("set_fake_env") +def test_deployment(watsonx_compat_mocks): + with requests_mock.Mocker() as m: + mock_response = watsonx_compat_mocks["watsonx_generation"] + extended_request = "/ml/v1/deployments/" + extended_request += "XXXXXXXX-XXXX-XXXX-XXXX-XXXXXXXXXXXX" + extended_request += "/text/generation?version=2021-05-01" + + m.post( + "https://garak.example.com/" + extended_request, json=mock_response["json"] + ) + + granite_llm = WatsonXGenerator(DEFAULT_DEPLOYMENT_NAME) + response = granite_llm._generate_with_deployment("What is this?") + + assert granite_llm.name == response["model_id"] diff --git a/tests/generators/watsonx.json b/tests/generators/watsonx.json new file mode 100644 index 000000000..6b1ef32ca --- /dev/null +++ b/tests/generators/watsonx.json @@ -0,0 +1,29 @@ +{ + "watsonx_bearer_token": { + "code": 200, + "json": { + "access_token": "fake_token1231231231", + "refresh_token": "not_supported", + "token_type": "Bearer", + "expires_in": 3600, + "expiration": 1737754747, + "scope": "ibm openid" + } + }, + "watsonx_generation": { + "code": 200, + "json" : { + "model_id": "ibm/granite-3-8b-instruct", + "model_version": "1.1.0", + "created_at": "2025-01-24T20:51:59.520Z", + "results": [ + { + "generated_text": "This is a test generation. :)", + "generated_token_count": 32, + "input_token_count": 6, + "stop_reason": "eos_token" + } + ] + } + } +} \ No newline at end of file From fefcd5e72726691e419d4fbceaa1c9bf98c2f931 Mon Sep 17 00:00:00 2001 From: iamnotcj <69222710+iamnotcj@users.noreply.github.com> Date: Thu, 30 Jan 2025 15:20:41 -0600 Subject: [PATCH 7/9] Update garak/generators/watsonx.py Added version parameter Co-authored-by: Jeffrey Martin Signed-off-by: iamnotcj <69222710+iamnotcj@users.noreply.github.com> --- garak/generators/watsonx.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/garak/generators/watsonx.py b/garak/generators/watsonx.py index 49b767dde..5c1b3b8b8 100644 --- a/garak/generators/watsonx.py +++ b/garak/generators/watsonx.py @@ -82,7 +82,7 @@ def _generate_with_deployment(self, payload): self.uri + "/ml/v1/deployments/" + self.deployment_id - + "/text/generation?version=2021-05-01" + + f"/text/generation?version={self.version}" ) body = {"parameters": {"prompt_variables": {self.prompt_variable: payload}}} headers = { From 4b252fea84876abdeb8dd0b7672c7037656a8d6b Mon Sep 17 00:00:00 2001 From: iamnotcj <69222710+iamnotcj@users.noreply.github.com> Date: Thu, 30 Jan 2025 15:21:25 -0600 Subject: [PATCH 8/9] Update garak/generators/watsonx.py Added version parameter Co-authored-by: Jeffrey Martin Signed-off-by: iamnotcj <69222710+iamnotcj@users.noreply.github.com> --- garak/generators/watsonx.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/garak/generators/watsonx.py b/garak/generators/watsonx.py index 5c1b3b8b8..f96449310 100644 --- a/garak/generators/watsonx.py +++ b/garak/generators/watsonx.py @@ -53,7 +53,7 @@ def _set_bearer_token(self, iam_url="https://iam.cloud.ibm.com/identity/token"): def _generate_with_project(self, payload): # Generation via Project ID. - url = self.uri + "/ml/v1/text/generation?version=2023-05-29" + url = self.uri + f"/ml/v1/text/generation?version={self.version}" body = { "input": payload, From 826f2dc15d1890b80a096be2df90a7c3aa831f29 Mon Sep 17 00:00:00 2001 From: iamnotcj <69222710+iamnotcj@users.noreply.github.com> Date: Thu, 30 Jan 2025 16:16:02 -0600 Subject: [PATCH 9/9] Updated version variable Updated the default version variable and updated tests. --- garak/generators/watsonx.py | 7 ++++--- tests/generators/test_watsonx.py | 2 +- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/garak/generators/watsonx.py b/garak/generators/watsonx.py index f96449310..e1e85bd3f 100644 --- a/garak/generators/watsonx.py +++ b/garak/generators/watsonx.py @@ -24,6 +24,7 @@ class WatsonXGenerator(Generator): DID_ENV_VAR = "WATSONX_DEPLOYID" DEFAULT_PARAMS = Generator.DEFAULT_PARAMS | { "uri": None, + "version": "2023-05-29", "project_id": "", "deployment_id": "", "prompt_variable": "input", @@ -38,7 +39,7 @@ def __init__(self, name="", config_root=_config): # Initialize and validate api_key if self.api_key is not None: os.environ[self.ENV_VAR] = self.api_key - + def _set_bearer_token(self, iam_url="https://iam.cloud.ibm.com/identity/token"): header = { "Content-Type": "application/x-www-form-urlencoded", @@ -52,7 +53,7 @@ def _set_bearer_token(self, iam_url="https://iam.cloud.ibm.com/identity/token"): def _generate_with_project(self, payload): # Generation via Project ID. - + url = self.uri + f"/ml/v1/text/generation?version={self.version}" body = { @@ -126,7 +127,7 @@ def _validate_env_var(self): def _call_model( self, prompt: str, generations_this_call: int = 1 ) -> List[Union[str, None]]: - if not self.bearer_token : + if not self.bearer_token: self._set_bearer_token() # Check if message is empty. If it is, append null byte. diff --git a/tests/generators/test_watsonx.py b/tests/generators/test_watsonx.py index e9683d559..4a2c2e16e 100644 --- a/tests/generators/test_watsonx.py +++ b/tests/generators/test_watsonx.py @@ -69,7 +69,7 @@ def test_deployment(watsonx_compat_mocks): mock_response = watsonx_compat_mocks["watsonx_generation"] extended_request = "/ml/v1/deployments/" extended_request += "XXXXXXXX-XXXX-XXXX-XXXX-XXXXXXXXXXXX" - extended_request += "/text/generation?version=2021-05-01" + extended_request += "/text/generation?version=2023-05-29" m.post( "https://garak.example.com/" + extended_request, json=mock_response["json"]