From 45087f32998aca37895aaaa3f5010391a7ec0391 Mon Sep 17 00:00:00 2001
From: Cyberes <cyberes@evulid.cc>
Date: Wed, 20 Sep 2023 15:24:44 -0600
Subject: [PATCH] check_vllm: fix

---
 check_vllm.py | 59 ++++++++++++++++++++++++++++-----------------------
 1 file changed, 33 insertions(+), 26 deletions(-)

diff --git a/check_vllm.py b/check_vllm.py
index 3ba3b04..55654e2 100644
--- a/check_vllm.py
+++ b/check_vllm.py
@@ -45,33 +45,40 @@ def main(critical, warning, timeout, target_url, does_not_contain, verify):
             }
         }
         if response.status_code != 200:
-            text_result = f"failed with status code {response.status_code}.\n{response.text}"
+            text_result = f"failed with status code {response.status_code}\n{response.text}"
             exit_code = nagios.STATE_CRIT
         else:
-            tokenizer = tiktoken.get_encoding("cl100k_base")
-            num_tokens = len(tokenizer.encode(json_data['text'][0]))
-            perfdata.update({
-                'tokens': {
-                    'value': num_tokens,
-                    'min': 0,
-                }
-            })
-
-            if not len(json_data['text'][0]):
-                text_result = f"response was empty.\n{response.text}"
-                exit_code = nagios.STATE_CRIT
-            elif elapsed_time >= warning:
-                text_result = f"response was {elapsed_time} seconds"
-                exit_code = nagios.STATE_WARN
-            elif elapsed_time >= critical:
-                text_result = f"response was {elapsed_time} seconds"
-                exit_code = nagios.STATE_CRIT
-            elif does_not_contain and does_not_contain in json_data['text'][0]:
-                text_result = f"\"{does_not_contain}\" was in the response:\n{json_data['text'][0]}"
+            if not json_data.get('text'):
+                text_result = f"did not receive valid JSON\n{response.text}"
                 exit_code = nagios.STATE_CRIT
             else:
-                text_result = f"generation time was {elapsed_time} seconds and generated {num_tokens} tokens"
-                exit_code = nagios.STATE_OK
+                if not len(json_data['text']):
+                    text_result = f"did not receive valid JSON\n{response.text}"
+                    exit_code = nagios.STATE_CRIT
+                else:
+                    tokenizer = tiktoken.get_encoding("cl100k_base")
+                    num_tokens = len(tokenizer.encode(json_data['text'][0]))
+                    perfdata.update({
+                        'tokens': {
+                            'value': num_tokens,
+                            'min': 0,
+                        }
+                    })
+                    if not len(json_data['text'][0]):
+                        text_result = f"response was empty.\n{response.text}"
+                        exit_code = nagios.STATE_CRIT
+                    elif elapsed_time >= warning:
+                        text_result = f"response was {elapsed_time} seconds"
+                        exit_code = nagios.STATE_WARN
+                    elif elapsed_time >= critical:
+                        text_result = f"response was {elapsed_time} seconds"
+                        exit_code = nagios.STATE_CRIT
+                    elif does_not_contain and does_not_contain in json_data['text'][0]:
+                        text_result = f"\"{does_not_contain}\" was in the response:\n{json_data['text'][0]}"
+                        exit_code = nagios.STATE_CRIT
+                    else:
+                        text_result = f"generation time was {elapsed_time} seconds and generated {num_tokens} tokens"
+                        exit_code = nagios.STATE_OK
     else:
         perfdata = {}
         text_result = f"CRITICAL - request failed - {error.__class__.__name__}: {error}"
@@ -87,12 +94,12 @@ if __name__ == "__main__":
     parser.add_argument('-w', '--warning', type=float, default=30, help='Warning threshold for response time in seconds (default: 30)')
     parser.add_argument('-t', '--target-url', type=str, required=True, help='Your completion API endpoint')
     parser.add_argument('-m', '--timeout', type=float, default=300, help='Request timeout in seconds (default: 300)')
-    parser.add_argument('-v', '--verify', action='store_false', help='Do not verify SSL')
-    parser.add_argument('-n', '--does-not-contain', type=str, help='Response must not contain')
+    parser.add_argument('-v', '--no-verify', action='store_false', help='Do not verify SSL')
+    parser.add_argument('-n', '--does-not-contain', type=str, help='Response must not contain this string')
     args = parser.parse_args()
 
     try:
-        main(args.critical, args.warning, args.timeout, args.target_url, args.does_not_contain, args.verify)
+        main(args.critical, args.warning, args.timeout, args.target_url, args.does_not_contain, args.no_verify)
     except Exception as e:
         print(f'UNKNOWN: exception "{e}"')
         print(traceback.format_exc())