File size: 5,129 Bytes
62274c6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | #!/usr/bin/env python3
"""Run bounded factual, JSON, code, tool, and optional vision checks."""
from __future__ import annotations
import argparse
import ast
import base64
import json
import time
from pathlib import Path
from common import endpoint, request_json, thinking_medium_sampler, write_report
def _message(response: dict) -> dict:
return response["choices"][0]["message"]
def validate(kind: str, response: dict) -> bool:
message = _message(response)
content = str(message.get("content") or "").strip()
if kind == "factual":
return content.rstrip(".") == "323"
if kind == "json":
return json.loads(content) == {"product": 323}
if kind == "code":
return _is_exact_square_function(content)
if kind == "tool":
calls = message.get("tool_calls") or []
return bool(calls and calls[0].get("function", {}).get("name") == "get_weather")
if kind == "vision":
return json.loads(content) == {"red_squares": 3, "blue_circles": 2}
raise ValueError(f"unknown fixture: {kind}")
def _is_exact_square_function(source: str) -> bool:
"""Validate the fixture's one safe function without executing model output."""
try:
tree = ast.parse(source)
except (SyntaxError, ValueError, TypeError):
return False
if len(tree.body) != 1 or not isinstance(tree.body[0], ast.FunctionDef):
return False
function = tree.body[0]
arguments = function.args
if (
function.name != "square"
or function.decorator_list
or function.returns is not None
or getattr(function, "type_params", ())
or arguments.posonlyargs
or len(arguments.args) != 1
or arguments.args[0].arg != "n"
or arguments.args[0].annotation is not None
or arguments.vararg is not None
or arguments.kwonlyargs
or arguments.kw_defaults
or arguments.kwarg is not None
or arguments.defaults
or len(function.body) != 1
or not isinstance(function.body[0], ast.Return)
):
return False
value = function.body[0].value
return (
isinstance(value, ast.BinOp)
and isinstance(value.op, ast.Mult)
and isinstance(value.left, ast.Name)
and value.left.id == "n"
and isinstance(value.left.ctx, ast.Load)
and isinstance(value.right, ast.Name)
and value.right.id == "n"
and isinstance(value.right.ctx, ast.Load)
)
def payload(kind: str, model: str, image: Path | None) -> dict[str, object]:
prompts = {
"factual": "Calculate 17 multiplied by 19. Return only the digits.",
"json": "Calculate 17 multiplied by 19. Return JSON exactly matching {\"product\": 323}.",
"code": "Return only Python source defining square(n) that returns n * n.",
"tool": "Use get_weather with city Boston. Do not answer from memory.",
"vision": "Count red squares and blue circles. Return JSON with those two keys.",
}
message: dict[str, object] = {"role": "user", "content": prompts[kind]}
if kind == "vision":
if image is None:
raise ValueError("--image is required with --include-vision")
message["content"] = [
{"type": "text", "text": prompts[kind]},
{"type": "image_url", "image_url": {"url": "data:image/png;base64," + base64.b64encode(image.read_bytes()).decode()}},
]
result: dict[str, object] = {"model": model, "messages": [message], "max_tokens": 512, **thinking_medium_sampler()}
if kind == "tool":
result["tools"] = [{"type": "function", "function": {"name": "get_weather", "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}}}]
if kind == "json":
result["response_format"] = {"type": "json_object"}
return result
def main(argv: list[str] | None = None) -> int:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--base-url", required=True)
parser.add_argument("--model", required=True)
parser.add_argument("--output", type=Path, required=True)
parser.add_argument("--timeout", type=float, default=1800.0)
parser.add_argument("--include-vision", action="store_true")
parser.add_argument("--image", type=Path)
args = parser.parse_args(argv)
if not 0 < args.timeout <= 1800:
parser.error("--timeout must be between 0 and 1800 seconds")
fixtures = ["factual", "json", "code", "tool"] + (["vision"] if args.include_vision else [])
results = {}
for kind in fixtures:
started = time.perf_counter()
response = request_json(endpoint(args.base_url, "chat/completions"), payload(kind, args.model, args.image), args.timeout)
results[kind] = {"passed": validate(kind, response), "elapsed_seconds": time.perf_counter() - started, "usage": response.get("usage", {})}
write_report(args.output, {"sampler": thinking_medium_sampler(), "results": results})
return 0 if all(item["passed"] for item in results.values()) else 1
if __name__ == "__main__":
raise SystemExit(main())
|