Add urllib.parse.urlparse to benchmark-local.py and validate that the provided --ollama-url uses either the http or https scheme. Fix arbitrary URI handling where the script could previously access local files (file://) or other unsupported protocols. If the scheme is invalid, parser.error is called to exit cleanly with a clear message.
This commit is contained in:
@@ -15,6 +15,7 @@ import json
|
||||
import re
|
||||
import time
|
||||
import urllib.request
|
||||
import urllib.parse
|
||||
from pathlib import Path
|
||||
|
||||
ROOT = Path(__file__).parent.parent
|
||||
@@ -149,6 +150,11 @@ def main():
|
||||
parser.add_argument("--repeat", type=int, default=1, help="Runs per cell; median reported (default: 1)")
|
||||
parser.add_argument("--ollama-url", default="http://localhost:11434", help="Ollama base URL")
|
||||
args = parser.parse_args()
|
||||
|
||||
parsed_url = urllib.parse.urlparse(args.ollama_url)
|
||||
if parsed_url.scheme not in ("http", "https"):
|
||||
parser.error(f"Invalid --ollama-url scheme: '{parsed_url.scheme}'. Only 'http' and 'https' are supported.")
|
||||
|
||||
run(args.model, args.repeat, args.ollama_url)
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user