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 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..e1e85bd3f --- /dev/null +++ b/garak/generators/watsonx.py @@ -0,0 +1,150 @@ +from garak import _config +from garak.generators.base import Generator +from typing import List, Union +import os +import requests + + +class WatsonXGenerator(Generator): + """ + This is a generator for watsonx.ai. + + 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" + URI_ENV_VAR = "WATSONX_URL" + PID_ENV_VAR = "WATSONX_PROJECTID" + DID_ENV_VAR = "WATSONX_DEPLOYID" + DEFAULT_PARAMS = Generator.DEFAULT_PARAMS | { + "uri": None, + "version": "2023-05-29", + "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 + f"/ml/v1/text/generation?version={self.version}" + + 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 + + f"/text/generation?version={self.version}" + ) + 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: + 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: + pass + else: + self.project_id = os.getenv("WATSONX_PROJECTID", "") + + # Initialize and validate deployment_id. + if self.deployment_id: + pass + else: + 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]]: + if not self.bearer_token: + self._set_bearer_token() + + # 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." + ) + + 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(output["results"][0]["generated_text"])] + + +DEFAULT_CLASS = "WatsonXGenerator" 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 74c2a153c..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 diff --git a/tests/generators/test_watsonx.py b/tests/generators/test_watsonx.py new file mode 100644 index 000000000..4a2c2e16e --- /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=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_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