Update handler.py

handler.py

@@ -1,38 +1,39 @@
 from typing import Dict, List, Any
 from llama_cpp import Llama
+import gemma_tools
 
 MAX_TOKENS = 8192
 
 
 class EndpointHandler():
     def __init__(self):
-        # Initialize the model with your ComicBot configuration
+        # Initialize the model with ComicBot configuration
         print("Initializing Llama model with ComicBot settings...")
         self.model = Llama.from_pretrained(
-            "njwright92/ComicBot_v.2-gguf", filename="comic_mistral-v5.2.q5_0.gguf", n_ctx=8192)
+            "njwright92/ComicBot_v.2-gguf", filename="comic_mistral-v5.2.q5_0.gguf", n_ctx=MAX_TOKENS)
         print("Model initialization complete.")
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
-        # Extract arguments from the data
-        print("Extracting arguments from the data payload...")
-        args = data.get("args", {})
-        print(f"Arguments extracted: {args}")
+        # Extract and validate arguments from the data
+        print("Extracting and validating arguments from the data payload...")
+        args_check = gemma_tools.get_args_or_none(
+            data)  # Using the new function
 
-        # Define the formatting template
-        fmat = "<startofturn>system\n{system_prompt} <endofturn>\n<startofturn>user\n{prompt} <endofturn>\n<startofturn>model"
-
-        # Check if args is properly formatted
-        if not args:
-            print("No arguments found in the data payload.")
+        if not args_check[0]:  # If validation failed
             return [{
-                "status": "error",
-                "message": "No arguments found in the data payload."
+                "status": args_check.get("status", "error"),
+                "reason": args_check.get("reason", "unknown"),
+                "description": args_check.get("description", "Validation error in arguments")
             }]
 
+        args = args_check  # If validation passed, args are in args_check
+
+        # Define the formatting template
+        fmat = "<startofturn>system\n{system_prompt} <endofturn>\n<startofturn>user\n{inputs} <endofturn>\n<startofturn>model"
+
         try:
-            fmat = fmat.format(system_prompt=args.get(
-                "system_prompt", ""), prompt=args.get("inputs", ""))
-            print(f"Formatted prompt: {fmat}")
+            formatted_prompt = fmat.format(**args)
+            print(f"Formatted prompt: {formatted_prompt}")
         except Exception as e:
             print(f"Error in formatting the prompt: {str(e)}")
             return [{
@@ -45,21 +46,24 @@ class EndpointHandler():
         try:
             max_length = int(max_length)
             print(f"Max length set to: {max_length}")
-        except Exception as e:
-            print(f"Error converting max_length to int: {str(e)}")
+        except ValueError:
             return [{
                 "status": "error",
-                "reason": "max_length was passed as something that was not a plain old int",
-                "detail": str(e)
+                "reason": "max_length must be an integer",
+                "detail": "max_length was not a valid integer"
             }]
 
         print("Generating response from the model...")
-        res = self.model(fmat, temperature=args.get("temperature", 1.0), top_p=args.get(
-            "top_p", 0.9), top_k=args.get("top_k", 40), max_tokens=max_length)
+        res = self.model(formatted_prompt,
+                         temperature=args["temperature"],
+                         top_p=args["top_p"],
+                         top_k=args["top_k"],
+                         max_tokens=max_length)
+
         print(f"Model response: {res}")
 
         return [{
             "status": "success",
-            "response": res
+            # Assuming Llama's response format
+            "response": res['choices'][0]['text']
         }]
-