Update vllm_support/src/kalm_t5gemma2_vllm_plugin/client.py

vllm_support/src/kalm_t5gemma2_vllm_plugin/client.py

@@ -86,17 +86,30 @@ def _load_payload(path: Path) -> dict[str, Any]:
     return payload
 
 
+def _request_payload(args: argparse.Namespace) -> dict[str, Any]:
+    payload = (
+        _load_payload(args.json_file)
+        if args.json_file is not None
+        else _demo_payload(args.endpoint, args.return_margin, args.top_k)
+    )
+    if args.return_margin:
+        payload["return_margin"] = True
+    if args.top_k is not None:
+        if args.endpoint != "rerank":
+            raise ValueError("--top-k is only valid with --endpoint rerank.")
+        if args.top_k < 0:
+            raise ValueError("--top-k must be non-negative.")
+        payload["top_k"] = args.top_k
+    return payload
+
+    
 def main() -> int:
     args = build_parser().parse_args()
     base_url = args.base_url.rstrip("/")
     if args.health:
         response = request_json("GET", f"{base_url}/health", timeout=args.timeout)
     else:
-        payload = (
-            _load_payload(args.json_file)
-            if args.json_file is not None
-            else _demo_payload(args.endpoint, args.return_margin, args.top_k)
-        )
+        payload = _request_payload(args)
         response = request_json(
             "POST",
             f"{base_url}/{args.endpoint}",
@@ -110,4 +123,3 @@ def main() -> int:
 
 if __name__ == "__main__":
     raise SystemExit(main())
-