From 3cc8fb2ab5f77aea255f7a179a6bffa4b18ff620 Mon Sep 17 00:00:00 2001 From: Leon Derczynski Date: Tue, 12 Nov 2024 15:38:17 +0100 Subject: [PATCH 1/8] add capability to skip generation for HTTP codes in given set --- garak/generators/rest.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/garak/generators/rest.py b/garak/generators/rest.py index 17ebe9b11..6c5506b30 100644 --- a/garak/generators/rest.py +++ b/garak/generators/rest.py @@ -29,7 +29,8 @@ class RestGenerator(Generator): DEFAULT_PARAMS = Generator.DEFAULT_PARAMS | { "headers": {}, "method": "post", - "ratelimit_codes": [429], + "ratelimit_codes": {429}, + "skip_codes": set(), "response_json": False, "response_json_field": None, "req_template": "$INPUT", @@ -194,7 +195,15 @@ def _call_model( } resp = self.http_function(self.uri, **req_kArgs) if resp.status_code in self.ratelimit_codes: - raise RateLimitHit(f"Rate limited: {resp.status_code} - {resp.reason}, uri: {self.uri}") + raise RateLimitHit( + f"Rate limited: {resp.status_code} - {resp.reason}, uri: {self.uri}" + ) + + elif resp.status_code in self.skip_codes: + logging.debug( + f"REST skip prompt: {resp.status_code} - {resp.reason}, uri: {self.uri}" + ) + return [None] elif str(resp.status_code)[0] == "3": raise NotImplementedError( From 0577ca0352da7e52b0600bc20279f0e69e8966bc Mon Sep 17 00:00:00 2001 From: Leon Derczynski Date: Tue, 12 Nov 2024 15:43:58 +0100 Subject: [PATCH 2/8] test RestGenerator skip code --- tests/generators/test_rest.py | 28 +++++++++++++++++++++++++++- 1 file changed, 27 insertions(+), 1 deletion(-) diff --git a/tests/generators/test_rest.py b/tests/generators/test_rest.py index 932473ba8..5c4913034 100644 --- a/tests/generators/test_rest.py +++ b/tests/generators/test_rest.py @@ -3,7 +3,7 @@ import requests_mock from sympy import is_increasing -from garak import _config +from garak import _config, _plugins from garak.generators.rest import RestGenerator @@ -95,3 +95,29 @@ def test_json_rest_deeper(requests_mock): generator = RestGenerator() output = generator._call_model("Who is Enabran Tain's son?") assert output == [DEFAULT_TEXT_RESPONSE] + + +@pytest.mark.usefixtures("set_rest_config") +def test_rest_skip_code(requests_mock): + generator = _plugins.load_plugin( + "generators.rest.RestGenerator", config_root=_config + ) + generator.skip_codes = {200} + requests_mock.post( + DEFAULT_URI, + text=json.dumps( + { + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": DEFAULT_TEXT_RESPONSE, + }, + } + ] + } + ), + ) + output = generator._call_model("Who is Enabran Tain's son?") + assert output == [None] From 645b585b06c7a3f4ccccb571086fa4e5816139b6 Mon Sep 17 00:00:00 2001 From: Leon Derczynski Date: Tue, 12 Nov 2024 15:47:51 +0100 Subject: [PATCH 3/8] add rest skip_codes doc; make skip_codes take precedence over ratelimit_codes --- docs/source/garak.generators.rest.rst | 1 + garak/generators/rest.py | 11 ++++++----- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/docs/source/garak.generators.rest.rst b/docs/source/garak.generators.rest.rst index 52130a566..5f0cf75bb 100644 --- a/docs/source/garak.generators.rest.rst +++ b/docs/source/garak.generators.rest.rst @@ -16,6 +16,7 @@ Uses the following options from ``_config.plugins.generators["rest.RestGenerator * ``response_json_field`` - (optional) Which field of the response JSON should be used as the output string? Default ``text``. Can also be a JSONPath value, and ``response_json_field`` is used as such if it starts with ``$``. * ``request_timeout`` - How many seconds should we wait before timing out? Default 20 * ``ratelimit_codes`` - Which endpoint HTTP response codes should be caught as indicative of rate limiting and retried? ``List[int]``, default ``[429]`` +* ``skip_codes`` - Which endpoint HTTP response code should lead to the generation being treated as not possible and skipped for this query. Takes precedence over ``skip_codes``. Templates can be either a string or a JSON-serialisable Python object. Instance of ``$INPUT`` here are replaced with the prompt; instances of ``$KEY`` diff --git a/garak/generators/rest.py b/garak/generators/rest.py index 6c5506b30..afc032dd8 100644 --- a/garak/generators/rest.py +++ b/garak/generators/rest.py @@ -194,17 +194,18 @@ def _call_model( "timeout": self.request_timeout, } resp = self.http_function(self.uri, **req_kArgs) - if resp.status_code in self.ratelimit_codes: - raise RateLimitHit( - f"Rate limited: {resp.status_code} - {resp.reason}, uri: {self.uri}" - ) - elif resp.status_code in self.skip_codes: + if resp.status_code in self.skip_codes: logging.debug( f"REST skip prompt: {resp.status_code} - {resp.reason}, uri: {self.uri}" ) return [None] + elif resp.status_code in self.ratelimit_codes: + raise RateLimitHit( + f"Rate limited: {resp.status_code} - {resp.reason}, uri: {self.uri}" + ) + elif str(resp.status_code)[0] == "3": raise NotImplementedError( f"REST URI redirection: {resp.status_code} - {resp.reason}, uri: {self.uri}" From e4078d76088e67dee046e075428a4946e3e1a17d Mon Sep 17 00:00:00 2001 From: Leon Derczynski Date: Tue, 12 Nov 2024 17:38:40 +0100 Subject: [PATCH 4/8] add skip_codes to _supported_params --- garak/generators/rest.py | 1 + 1 file changed, 1 insertion(+) diff --git a/garak/generators/rest.py b/garak/generators/rest.py index afc032dd8..095f4e230 100644 --- a/garak/generators/rest.py +++ b/garak/generators/rest.py @@ -56,6 +56,7 @@ class RestGenerator(Generator): "req_template_json_object", "request_timeout", "ratelimit_codes", + "skip_codes", "temperature", "top_k", ) From 2b5f109a4224c07adefd0652b19903f4468bd83e Mon Sep 17 00:00:00 2001 From: Leon Derczynski Date: Tue, 12 Nov 2024 21:08:22 +0100 Subject: [PATCH 5/8] typo in docs/source/garak.generators.rest.rst Co-authored-by: Jeffrey Martin Signed-off-by: Leon Derczynski --- docs/source/garak.generators.rest.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/garak.generators.rest.rst b/docs/source/garak.generators.rest.rst index 5f0cf75bb..6d303e063 100644 --- a/docs/source/garak.generators.rest.rst +++ b/docs/source/garak.generators.rest.rst @@ -16,7 +16,7 @@ Uses the following options from ``_config.plugins.generators["rest.RestGenerator * ``response_json_field`` - (optional) Which field of the response JSON should be used as the output string? Default ``text``. Can also be a JSONPath value, and ``response_json_field`` is used as such if it starts with ``$``. * ``request_timeout`` - How many seconds should we wait before timing out? Default 20 * ``ratelimit_codes`` - Which endpoint HTTP response codes should be caught as indicative of rate limiting and retried? ``List[int]``, default ``[429]`` -* ``skip_codes`` - Which endpoint HTTP response code should lead to the generation being treated as not possible and skipped for this query. Takes precedence over ``skip_codes``. +* ``skip_codes`` - Which endpoint HTTP response code should lead to the generation being treated as not possible and skipped for this query. Takes precedence over ``ratelimit_codes``. Templates can be either a string or a JSON-serialisable Python object. Instance of ``$INPUT`` here are replaced with the prompt; instances of ``$KEY`` From 2983ac965bb9ed72780a3399407dd8c2ecbe18e9 Mon Sep 17 00:00:00 2001 From: Leon Derczynski Date: Tue, 12 Nov 2024 21:20:38 +0100 Subject: [PATCH 6/8] type skip_codes, retry_codes as lists; satisfy reasonable linter requests --- garak/generators/rest.py | 24 ++++++++++++++---------- 1 file changed, 14 insertions(+), 10 deletions(-) diff --git a/garak/generators/rest.py b/garak/generators/rest.py index 095f4e230..cd4eca2b9 100644 --- a/garak/generators/rest.py +++ b/garak/generators/rest.py @@ -29,8 +29,8 @@ class RestGenerator(Generator): DEFAULT_PARAMS = Generator.DEFAULT_PARAMS | { "headers": {}, "method": "post", - "ratelimit_codes": {429}, - "skip_codes": set(), + "ratelimit_codes": [429], + "skip_codes": [], "response_json": False, "response_json_field": None, "req_template": "$INPUT", @@ -123,7 +123,7 @@ def __init__(self, uri=None, config_root=_config): try: self.json_expr = jsonpath_ng.parse(self.response_json_field) except JsonPathParserError as e: - logging.CRITICAL( + logging.critical( "Couldn't parse response_json_field %s", self.response_json_field ) raise e @@ -198,37 +198,41 @@ def _call_model( if resp.status_code in self.skip_codes: logging.debug( - f"REST skip prompt: {resp.status_code} - {resp.reason}, uri: {self.uri}" + "REST skip prompt: %s - %s, uri: %s", + resp.status_code, + resp.reason, + self.uri, ) return [None] - elif resp.status_code in self.ratelimit_codes: + if resp.status_code in self.ratelimit_codes: raise RateLimitHit( f"Rate limited: {resp.status_code} - {resp.reason}, uri: {self.uri}" ) - elif str(resp.status_code)[0] == "3": + if str(resp.status_code)[0] == "3": raise NotImplementedError( f"REST URI redirection: {resp.status_code} - {resp.reason}, uri: {self.uri}" ) - elif str(resp.status_code)[0] == "4": + if str(resp.status_code)[0] == "4": raise ConnectionError( f"REST URI client error: {resp.status_code} - {resp.reason}, uri: {self.uri}" ) - elif str(resp.status_code)[0] == "5": + if str(resp.status_code)[0] == "5": error_msg = f"REST URI server error: {resp.status_code} - {resp.reason}, uri: {self.uri}" if self.retry_5xx: raise IOError(error_msg) - else: - raise ConnectionError(error_msg) + raise ConnectionError(error_msg) if not self.response_json: return [str(resp.text)] response_object = json.loads(resp.content) + response = [None] * generations_this_call + # if response_json_field starts with a $, treat is as a JSONPath assert ( self.response_json From bc5fb2439164c81ed3df8f767e63c16f0e6d77db Mon Sep 17 00:00:00 2001 From: Leon Derczynski Date: Thu, 14 Nov 2024 10:39:58 +0100 Subject: [PATCH 7/8] move set to list param Co-authored-by: Jeffrey Martin Signed-off-by: Leon Derczynski --- tests/generators/test_rest.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/generators/test_rest.py b/tests/generators/test_rest.py index 5c4913034..6f0949def 100644 --- a/tests/generators/test_rest.py +++ b/tests/generators/test_rest.py @@ -102,7 +102,7 @@ def test_rest_skip_code(requests_mock): generator = _plugins.load_plugin( "generators.rest.RestGenerator", config_root=_config ) - generator.skip_codes = {200} + generator.skip_codes = [200] requests_mock.post( DEFAULT_URI, text=json.dumps( From effa9436aee07643e43169e34bf32e4559bc98d4 Mon Sep 17 00:00:00 2001 From: Leon Derczynski Date: Thu, 14 Nov 2024 10:41:40 +0100 Subject: [PATCH 8/8] set default restgenerator response to [None] (i.e. generations_this_call == 1) Co-authored-by: Jeffrey Martin Signed-off-by: Leon Derczynski --- garak/generators/rest.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/garak/generators/rest.py b/garak/generators/rest.py index cd4eca2b9..5dfa6b273 100644 --- a/garak/generators/rest.py +++ b/garak/generators/rest.py @@ -231,7 +231,7 @@ def _call_model( response_object = json.loads(resp.content) - response = [None] * generations_this_call + response = [None] # if response_json_field starts with a $, treat is as a JSONPath assert (