#!/usr/bin/env python3

"""

Test structured output capability via chat completions endpoint.

Each test case contains:

- response_format: OpenAI-compatible response_format specification.

Both "json_schema" and "json_object" are accepted; with

"json_object" a schema can be supplied via extra_body.

- extra_body (optional): dict of extra top-level request fields merged into

the request payload (mirrors the OpenAI SDK's extra_body

feature; llama.cpp reads a top-level "json_schema" here).

- messages: initial conversation messages

- tools (optional): tool definitions (for mixed tool + structured tests)

- mock_tool_responses (optional): dict mapping tool_name -> callable(arguments) -> str (JSON)

- apply_stage: "always" to apply response_format to every request,

"after_tools" to run the tool loop plain, then request a

structured summary in a follow-up user turn.

- followup (optional, for after_tools): user message appended before the

final structured call.

- validate: callable(parsed_json, tool_calls_history, raw_content) -> (passed: bool, reason: str)

"""

import argparse

import json

import requests

import sys

from typing import Any, cast

# ---------------------------------------------------------------------------

# Color / formatting helpers

# ---------------------------------------------------------------------------

RESET = "\x1b[0m"

BOLD = "\x1b[1m"

DIM = "\x1b[2m"

CYAN = "\x1b[36m"

YELLOW = "\x1b[33m"

GREEN = "\x1b[32m"

RED = "\x1b[31m"

BLUE = "\x1b[34m"

WHITE = "\x1b[97m"

MAGENTA = "\x1b[35m"

def _print(text="", end="\n"):

sys.stdout.write(text + end)

sys.stdout.flush()

def print_header(title):

bar = "─" * 60

_print(f"\n{BOLD}{CYAN}┌{bar}┐{RESET}")

_print(

f"{BOLD}{CYAN}│ {WHITE}{title}{CYAN}{' ' * max(0, 58 - len(title))}│{RESET}"

)

_print(f"{BOLD}{CYAN}└{bar}┘{RESET}")

def print_tool_call(name, args):

args_str = json.dumps(args)

_print(

f"\n {BOLD}{YELLOW}⚙ tool call{RESET} {CYAN}{name}{RESET}{DIM}({args_str}){RESET}"

)

def print_tool_result(result):

preview = result[:160] + ("…" if len(result) > 160 else "")

_print(f" {DIM}{BLUE}↳ result{RESET} {DIM}{preview}{RESET}")

def print_model_output(text):

sys.stdout.write(text)

sys.stdout.flush()

def print_pass(reason):

_print(f"\n{BOLD}{GREEN}✔ PASS{RESET} {reason}")

def print_fail(reason):

_print(f"\n{BOLD}{RED}✘ FAIL{RESET} {reason}")

def print_info(msg):

_print(f"{DIM}{msg}{RESET}")

def print_schema_note(label, rf, extra_body=None):

kind = rf.get("type", "?")

name = ""

if kind == "json_schema":

name = rf.get("json_schema", {}).get("name", "")

elif kind == "json_object" and extra_body and "json_schema" in extra_body:

extra_schema = extra_body["json_schema"] or {}

name = extra_schema.get("title") or "extra_body.json_schema"

_print(f"{DIM}{MAGENTA} ⟐ response_format [{label}]: {kind}"

f"{(' / ' + name) if name else ''}{RESET}")

# ---------------------------------------------------------------------------

# HTTP helpers

# ---------------------------------------------------------------------------

def chat_completion(url, messages, tools=None, response_format=None, stream=False,

extra_body=None):

payload = {

"messages": messages,

"stream": stream,

"max_tokens": 8192,

}

if tools:

payload["tools"] = tools

payload["tool_choice"] = "auto"

if response_format is not None:

payload["response_format"] = response_format

if extra_body:

payload.update(extra_body)

try:

response = requests.post(url, json=payload, stream=stream)

response.raise_for_status()

except requests.exceptions.RequestException as e:

body = e.response.content if (e.response is not None) else b""

print_fail(f"Request error: {e} | body: {body}")

return None

full_content = ""

reasoning_content = ""

tool_calls: list[dict] = []

if stream:

for line in response.iter_lines():

if not line:

continue

decoded = line.decode("utf-8")

if not decoded.startswith("data: "):

continue

data_str = decoded[6:]

if data_str == "[DONE]":

break

try:

data = json.loads(data_str)

except json.JSONDecodeError:

continue

choices = data.get("choices", [])

if not choices:

continue

delta = choices[0].get("delta", {})

if delta.get("reasoning_content"):

reasoning_content += delta["reasoning_content"]

if delta.get("content"):

full_content += delta["content"]

print_model_output(delta["content"])

for tc in delta.get("tool_calls", []):

idx = tc.get("index", 0)

while len(tool_calls) <= idx:

tool_calls.append(

{

"id": "",

"type": "function",

"function": {"name": "", "arguments": ""},

}

)

if "id" in tc:

tool_calls[idx]["id"] += tc["id"]

if "function" in tc:

if "name" in tc["function"]:

tool_calls[idx]["function"]["name"] += tc["function"]["name"]

if "arguments" in tc["function"]:

tool_calls[idx]["function"]["arguments"] += tc["function"][

"arguments"

]

else:

data = response.json()

choices = data.get("choices", [])

if choices:

msg = choices[0].get("message", {})

full_content = msg.get("content") or ""

reasoning_content = msg.get("reasoning_content") or ""

tool_calls = msg.get("tool_calls") or []

if full_content:

print_model_output(full_content)

result = {"content": full_content, "tool_calls": tool_calls}

if reasoning_content:

result["reasoning_content"] = reasoning_content

return result

def run_tool_loop(

url, messages, tools, mock_tool_responses, stream, response_format=None,

extra_body=None, max_turns=6,

):

"""

Drive the tool-call loop. If response_format is provided it is applied to

every request. Returns (all_tool_calls, final_messages, final_content).

"""

msgs = list(messages)

all_tool_calls: list[dict] = []

for _ in range(max_turns):

result = chat_completion(

url, msgs, tools=tools, response_format=response_format, stream=stream,

extra_body=extra_body,

)

if result is None:

return all_tool_calls, msgs, None

tcs = result.get("tool_calls") or []

content = result.get("content") or ""

if not tcs:

if content:

_print(f"\n{DIM}{'·' * 60}{RESET}")

return all_tool_calls, msgs, content

all_tool_calls.extend(tcs)

assistant_msg: dict = {

"role": "assistant",

"content": content,

"tool_calls": tcs,

}

reasoning = result.get("reasoning_content")

if reasoning:

assistant_msg["reasoning_content"] = reasoning

msgs.append(assistant_msg)

for tc in tcs:

tool_name = tc["function"]["name"]

try:

args = json.loads(tc["function"]["arguments"])

except json.JSONDecodeError:

args = {}

print_tool_call(tool_name, args)

mock_fn = mock_tool_responses.get(tool_name) if mock_tool_responses else None

if mock_fn:

tool_result = mock_fn(args)

else:

tool_result = json.dumps({"error": f"Unknown tool: {tool_name}"})

print_tool_result(tool_result)

msgs.append(

{

"role": "tool",

"tool_call_id": tc.get("id", ""),

"content": tool_result,

}

)

return all_tool_calls, msgs, None

# ---------------------------------------------------------------------------

# Test case runner

# ---------------------------------------------------------------------------

def _try_parse_json(text):

"""Attempt to parse text as JSON, trimming common markdown fences."""

if text is None:

return None

stripped = text.strip()

if stripped.startswith("```"):

lines = stripped.splitlines()

if lines and lines[0].startswith("```"):

lines = lines[1:]

if lines and lines[-1].strip().startswith("```"):

lines = lines[:-1]

stripped = "\n".join(lines).strip()

try:

return json.loads(stripped)

except json.JSONDecodeError:

return None

def run_test(url, test_case, stream):

name = test_case["name"]

mode = f"{'stream' if stream else 'non-stream'}"

apply_stage = test_case.get("apply_stage", "always")

print_header(f"{name} [{mode}] ({apply_stage})")

response_format = test_case["response_format"]

extra_body = test_case.get("extra_body")

print_schema_note(apply_stage, response_format, extra_body)

tools = test_case.get("tools")

mocks = test_case.get("mock_tool_responses") or {}

all_tcs: list[dict] = []

final_content = None

if apply_stage == "always":

all_tcs, _msgs, final_content = run_tool_loop(

url,

messages=list(test_case["messages"]),

tools=tools,

mock_tool_responses=mocks,

stream=stream,

response_format=response_format,

extra_body=extra_body,

)

elif apply_stage == "after_tools":

# Phase 1: plain tool loop, no response_format applied yet.

all_tcs, msgs, interim_content = run_tool_loop(

url,

messages=list(test_case["messages"]),

tools=tools,

mock_tool_responses=mocks,

stream=stream,

response_format=None,

)

if interim_content:

msgs.append({"role": "assistant", "content": interim_content})

followup = test_case.get(

"followup",

"Now output the answer strictly as JSON matching the provided schema. "

"Do not include commentary.",

)

msgs.append({"role": "user", "content": followup})

# Phase 2: request final structured output. Tools are not passed so the

# model focuses on producing the schema-constrained answer.

_print(f"\n{DIM}{MAGENTA} ⟐ follow-up turn with response_format applied{RESET}")

result = chat_completion(

url, msgs, tools=None, response_format=response_format, stream=stream,

extra_body=extra_body,

)

final_content = result["content"] if result else None

else:

print_fail(f"Unknown apply_stage: {apply_stage}")

return False

if final_content is None:

print_fail("No final content from server.")

return False

parsed = _try_parse_json(final_content)

if parsed is None:

print_fail(f"Final content is not valid JSON: {final_content[:200]!r}")

return False

passed, reason = test_case["validate"](parsed, all_tcs, final_content)

if passed:

print_pass(reason)

else:

print_fail(reason)

return passed

# ---------------------------------------------------------------------------

# Test case definitions

# ---------------------------------------------------------------------------

# ---- Test 1: Book metadata extraction (always / json_schema) ----

_BOOK_SCHEMA = {

"type": "json_schema",

"json_schema": {

"name": "book_metadata",

"strict": True,

"schema": {

"type": "object",

"additionalProperties": False,

"properties": {

"title": {"type": "string"},

"author": {"type": "string"},

"year": {"type": "integer"},

"genre": {

"type": "string",

"enum": [

"fiction",

"non-fiction",

"fantasy",

"sci-fi",

"mystery",

"biography",

"history",

"other",

],

},

"page_count": {"type": "integer"},

},

"required": ["title", "author", "year", "genre", "page_count"],

},

},

}

BOOK_TEST_CASE = {

"name": "Book metadata extraction (json_schema, always)",

"response_format": _BOOK_SCHEMA,

"apply_stage": "always",

"messages": [

{

"role": "user",

"content": (

"Extract book metadata from this description: "

"'Dune is a 1965 science fiction epic by Frank Herbert, spanning roughly "

"688 pages in its first edition, set on the desert planet Arrakis.' "

"Return the data as JSON."

),

}

],

"validate": lambda parsed, tcs, raw: _validate_book(parsed),

}

def _validate_book(parsed):

required = {"title", "author", "year", "genre", "page_count"}

missing = required - parsed.keys()

if missing:

return False, f"Missing fields: {missing}"

if not isinstance(parsed["title"], str) or not parsed["title"]:

return False, "title must be a non-empty string"

if not isinstance(parsed["author"], str) or "herbert" not in parsed["author"].lower():

return False, f"author unexpected: {parsed['author']!r}"

if not isinstance(parsed["year"], int) or parsed["year"] != 1965:

return False, f"year should be 1965, got {parsed['year']!r}"

if parsed["genre"] not in {

"fiction", "non-fiction", "fantasy", "sci-fi", "mystery",

"biography", "history", "other",

}:

return False, f"genre not in enum: {parsed['genre']!r}"

if not isinstance(parsed["page_count"], int) or parsed["page_count"] <= 0:

return False, f"page_count should be positive int: {parsed['page_count']!r}"

return True, f"Book: {parsed['title']} ({parsed['year']}) / {parsed['genre']}"

# ---- Test 2: Sentiment classification (always / enum-constrained) ----

_SENTIMENT_SCHEMA = {

"type": "json_schema",

"json_schema": {

"name": "sentiment_analysis",

"strict": True,

"schema": {

"type": "object",

"additionalProperties": False,

"properties": {

"sentiment": {

"type": "string",

"enum": ["positive", "negative", "neutral"],

},

"confidence": {"type": "number"},

"keywords": {

"type": "array",

"items": {"type": "string"},

"minItems": 1,

"maxItems": 5,

},

},

"required": ["sentiment", "confidence", "keywords"],

},

},

}

SENTIMENT_TEST_CASE = {

"name": "Sentiment analysis with enum and array",

"response_format": _SENTIMENT_SCHEMA,

"apply_stage": "always",

"messages": [

{

"role": "user",

"content": (

"Analyse the sentiment of this review and return JSON with the "

"detected sentiment label, a confidence score between 0 and 1, "

"and up to five keyword strings that drove the classification:\n\n"

"'This product completely exceeded my expectations. The build "

"quality is phenomenal, it arrived a day early, and customer "

"support was delightful when I had a setup question.'"

),

}

],

"validate": lambda parsed, tcs, raw: _validate_sentiment(parsed),

}

def _validate_sentiment(parsed):

if parsed.get("sentiment") not in {"positive", "negative", "neutral"}:

return False, f"sentiment not in enum: {parsed.get('sentiment')!r}"

if parsed["sentiment"] != "positive":

return False, f"expected positive sentiment, got {parsed['sentiment']}"

conf = parsed.get("confidence")

if not isinstance(conf, (int, float)) or not (0.0 <= conf <= 1.0):

return False, f"confidence not in [0,1]: {conf!r}"

kws = parsed.get("keywords")

if not isinstance(kws, list) or not (1 <= len(kws) <= 5):

return False, f"keywords length out of range: {kws!r}"

if not all(isinstance(k, str) and k for k in kws):

return False, f"keywords must be non-empty strings: {kws!r}"

return True, f"sentiment={parsed['sentiment']} conf={conf} kws={kws}"

# ---- Test: json_object + extra_body.json_schema (always) ----

#

# Exercises the llama.cpp-specific path where the OpenAI SDK would send

# response_format={"type": "json_object"} and tunnel the schema through

# extra_body.json_schema (which becomes a top-level "json_schema" field on

# the request body).

_PRODUCT_JSON_OBJECT_SCHEMA = {

"$schema": "https://json-schema.org/draft/2020-12/schema",

"$id": "https://example.com/product.schema.json",

"title": "Product",

"description": "A product in the catalog",

"type": "object",

}

PRODUCT_JSON_OBJECT_TEST_CASE = {

"name": "json_object response_format with extra_body json_schema",

"response_format": {"type": "json_object"},

"extra_body": {"json_schema": _PRODUCT_JSON_OBJECT_SCHEMA},

"apply_stage": "always",

"messages": [

{

"role": "system",

"content": (

"Extract structured data from the provided text according to the "

"JSON schema. Return only valid JSON matching the schema exactly."

),

},

{

"role": "user",

"content": "Product: Wireless Headphones, ID: 101, In Stock: Yes",

},

],

"validate": lambda parsed, tcs, raw: _validate_product_json_object(parsed),

}

def _validate_product_json_object(parsed):

if not isinstance(parsed, dict):

return False, f"expected JSON object, got {type(parsed).__name__}: {parsed!r}"

if not parsed:

return False, f"expected non-empty object, got {parsed!r}"

return True, f"product object with {len(parsed)} field(s): {sorted(parsed.keys())}"

# ---- Test 3: Nested recipe schema (always) ----

_RECIPE_SCHEMA = {

"type": "json_schema",

"json_schema": {

"name": "recipe",

"strict": True,

"schema": {

"type": "object",

"additionalProperties": False,

"properties": {

"name": {"type": "string"},

"servings": {"type": "integer"},

"ingredients": {

"type": "array",

"minItems": 2,

"items": {

"type": "object",

"additionalProperties": False,

"properties": {

"item": {"type": "string"},

"quantity": {"type": "string"},

},

"required": ["item", "quantity"],

},

},

"steps": {

"type": "array",

"minItems": 2,

"items": {"type": "string"},

},

"prep_time_minutes": {"type": "integer"},

},

"required": ["name", "servings", "ingredients", "steps", "prep_time_minutes"],

},

},

}

RECIPE_TEST_CASE = {

"name": "Nested recipe with arrays of objects",

"response_format": _RECIPE_SCHEMA,

"apply_stage": "always",

"messages": [

{

"role": "user",

"content": (

"Give me a simple 4-serving scrambled eggs recipe as structured JSON. "

"Include the recipe name, servings, ingredients (each with item and "

"quantity), preparation steps, and total prep time in minutes."

),

}

],

"validate": lambda parsed, tcs, raw: _validate_recipe(parsed),

}

def _validate_recipe(parsed):

required = {"name", "servings", "ingredients", "steps", "prep_time_minutes"}

missing = required - parsed.keys()

if missing:

return False, f"Missing fields: {missing}"

if not isinstance(parsed["name"], str) or not parsed["name"]:

return False, "name must be a non-empty string"

if not isinstance(parsed["servings"], int) or parsed["servings"] <= 0:

return False, f"servings must be positive int: {parsed['servings']!r}"

ings = parsed["ingredients"]

if not isinstance(ings, list) or len(ings) < 2:

return False, f"ingredients must be array of >=2: got {ings!r}"

for i, ing in enumerate(ings):

if not isinstance(ing, dict):

return False, f"ingredient[{i}] is not an object: {ing!r}"

ing_d = cast(dict[str, Any], ing)

item_val = ing_d.get("item")

qty_val = ing_d.get("quantity")

if item_val is None or qty_val is None:

return False, f"ingredient[{i}] missing item/quantity: {ing!r}"

if not isinstance(item_val, str) or not isinstance(qty_val, str):

return False, f"ingredient[{i}] fields must be strings: {ing!r}"

steps = parsed["steps"]

if not isinstance(steps, list) or len(steps) < 2:

return False, f"steps must be array of >=2 strings: got {steps!r}"

if not all(isinstance(s, str) and s for s in steps):

return False, "all steps must be non-empty strings"

pt = parsed["prep_time_minutes"]

if not isinstance(pt, int) or pt <= 0:

return False, f"prep_time_minutes must be positive int: {pt!r}"

return True, f"recipe '{parsed['name']}' with {len(ings)} ingredients, {len(steps)} steps"

# ---- Test 4: Tool call -> structured product comparison (after_tools) ----

_SHOP_TOOLS = [

{

"type": "function",

"function": {

"name": "search_products",

"description": "Search a product catalogue by keyword.",

"parameters": {

"type": "object",

"properties": {

"query": {"type": "string"},

},

"required": ["query"],

},

},

},

{

"type": "function",

"function": {

"name": "get_product_details",

"description": "Get detailed specs for a product by ID.",

"parameters": {

"type": "object",

"properties": {

"product_id": {"type": "string"},

},

"required": ["product_id"],

},

},

},

]

_SHOP_SEARCH_RESULT = {

"results": [

{"product_id": "LAP-001", "title": "AeroBook 13 Pro", "price": 1399.0, "rating": 4.7},

{"product_id": "LAP-002", "title": "QuantumSlim 14", "price": 1199.0, "rating": 4.4},

{"product_id": "LAP-003", "title": "NimbusWork Ultra 15", "price": 999.0, "rating": 4.2},

],

}

_SHOP_PRODUCT_DETAILS = {

"LAP-001": {

"product_id": "LAP-001",

"title": "AeroBook 13 Pro",

"cpu": "M-series 10-core",

"ram_gb": 16,

"storage_gb": 512,

"battery_hours": 18,

"weight_kg": 1.24,

"price": 1399.0,

},

"LAP-002": {

"product_id": "LAP-002",

"title": "QuantumSlim 14",

"cpu": "Core i7 12-core",

"ram_gb": 16,

"storage_gb": 512,

"battery_hours": 12,

"weight_kg": 1.35,

"price": 1199.0,

},

"LAP-003": {

"product_id": "LAP-003",

"title": "NimbusWork Ultra 15",

"cpu": "Ryzen 7 8-core",

"ram_gb": 16,

"storage_gb": 1024,

"battery_hours": 10,

"weight_kg": 1.70,

"price": 999.0,

},

}

def _shop_details_mock(args):

pid = args.get("product_id", "")

if pid in _SHOP_PRODUCT_DETAILS:

return json.dumps(_SHOP_PRODUCT_DETAILS[pid])

return json.dumps({"error": f"unknown product_id: {pid}"})

_SHOP_COMPARISON_SCHEMA = {

"type": "json_schema",

"json_schema": {

"name": "laptop_comparison",

"strict": True,

"schema": {

"type": "object",

"additionalProperties": False,

"properties": {

"recommendation": {"type": "string"},

"ranked_candidates": {

"type": "array",

"minItems": 2,

"items": {

"type": "object",

"additionalProperties": False,

"properties": {

"product_id": {"type": "string"},

"title": {"type": "string"},

"score": {"type": "number"},

"reason": {"type": "string"},

},

"required": ["product_id", "title", "score", "reason"],

},

},

},

"required": ["recommendation", "ranked_candidates"],

},

},

}

SHOP_COMPARISON_TEST_CASE = {

"name": "Tool calls then structured laptop comparison (after_tools)",

"response_format": _SHOP_COMPARISON_SCHEMA,

"apply_stage": "after_tools",

"tools": _SHOP_TOOLS,

"mock_tool_responses": {

"search_products": lambda _: json.dumps(_SHOP_SEARCH_RESULT),

"get_product_details": _shop_details_mock,

},

"messages": [

{

"role": "user",

"content": (

"I need a lightweight laptop for travel. Please search the catalogue "

"for 'ultraportable laptop', then fetch detailed specs for at least two "

"of the top candidates. Once you've gathered the data I'll ask you to "

"produce a structured comparison."

),

}

],

"followup": (

"Thanks. Now produce the final comparison strictly as JSON matching the "

"laptop_comparison schema: your single best recommendation (the product_id), "

"and a ranked_candidates array of at least two laptops, each with "

"product_id, title, a numeric score, and a short reason."

),

"validate": lambda parsed, tcs, raw: _validate_shop_comparison(parsed, tcs),

}

def _validate_shop_comparison(parsed, tcs):

names = [tc["function"]["name"] for tc in tcs]

if "search_products" not in names:

return False, f"expected search_products tool call, got {names}"

if "get_product_details" not in names:

return False, f"expected get_product_details tool call, got {names}"

if "recommendation" not in parsed or not isinstance(parsed["recommendation"], str):

return False, f"recommendation missing or not a string: {parsed!r}"

cands = parsed.get("ranked_candidates")

if not isinstance(cands, list) or len(cands) < 2:

return False, f"ranked_candidates must be >=2: {cands!r}"

valid_ids = set(_SHOP_PRODUCT_DETAILS.keys())

candidate_pids: list = []

for i, c in enumerate(cands):

if not isinstance(c, dict):

return False, f"candidate[{i}] not an object: {c!r}"

c_d = cast(dict[str, Any], c)

pid = c_d.get("product_id")

title = c_d.get("title")

score = c_d.get("score")

reason = c_d.get("reason")

for k, v in (("product_id", pid), ("title", title),

("score", score), ("reason", reason)):

if v is None:

return False, f"candidate[{i}] missing {k}: {c!r}"

if pid not in valid_ids:

return False, f"candidate[{i}].product_id not in catalogue: {pid!r}"

if not isinstance(score, (int, float)):

return False, f"candidate[{i}].score not numeric: {score!r}"

candidate_pids.append(pid)

recommendation = parsed["recommendation"]

if recommendation not in valid_ids and recommendation not in candidate_pids:

return False, f"recommendation {recommendation!r} not in candidates"

return True, (

f"tools={names}; recommended={parsed['recommendation']}; "

f"{len(cands)} ranked candidates"

)

# ---- Test 5: Multi-step research then structured report (after_tools) ----

_RESEARCH_TOOLS = [

{

"type": "function",

"function": {

"name": "get_country_stats",

"description": "Fetch basic statistics for a country (population, GDP, capital).",

"parameters": {

"type": "object",

"properties": {

"country": {"type": "string"},

},

"required": ["country"],

},

},

},

{

"type": "function",

"function": {

"name": "get_climate_info",

"description": "Fetch climate information for a country.",

"parameters": {

"type": "object",

"properties": {

"country": {"type": "string"},

},

"required": ["country"],

},

},

},

]

_COUNTRY_STATS = {

"norway": {

"country": "Norway",

"capital": "Oslo",

"population": 5_480_000,

"gdp_usd_trillion": 0.48,

"currency": "NOK",

}

}

_CLIMATE_INFO = {

"norway": {

"country": "Norway",

"climate_zone": "subarctic / temperate coastal",

"avg_winter_temp_c": -4.5,

"avg_summer_temp_c": 16.0,

"annual_precipitation_mm": 1400,

}

}

def _country_stats_mock(args):

c = args.get("country", "").strip().lower()

if c in _COUNTRY_STATS:

return json.dumps(_COUNTRY_STATS[c])

return json.dumps({"error": f"unknown country: {c}"})

def _climate_info_mock(args):

c = args.get("country", "").strip().lower()

if c in _CLIMATE_INFO:

return json.dumps(_CLIMATE_INFO[c])

return json.dumps({"error": f"unknown country: {c}"})

_RESEARCH_REPORT_SCHEMA = {

"type": "json_schema",

"json_schema": {

"name": "country_report",

"strict": True,

"schema": {

"type": "object",

"additionalProperties": False,

"properties": {

"country": {"type": "string"},

"capital": {"type": "string"},

"population": {"type": "integer"},

"climate_summary": {"type": "string"},

"highlights": {

"type": "array",

"minItems": 2,

"maxItems": 5,

"items": {"type": "string"},

},

"suitable_for_tourism": {"type": "boolean"},

},

"required": [

"country", "capital", "population",

"climate_summary", "highlights", "suitable_for_tourism",

],

},

},

}

COUNTRY_REPORT_TEST_CASE = {

"name": "Research pipeline then structured country report (after_tools)",

"response_format": _RESEARCH_REPORT_SCHEMA,

"apply_stage": "after_tools",

"tools": _RESEARCH_TOOLS,

"mock_tool_responses": {

"get_country_stats": _country_stats_mock,

"get_climate_info": _climate_info_mock,

},

"messages": [

{

"role": "user",

"content": (

"I'm preparing a short briefing on Norway. Please call the "

"get_country_stats and get_climate_info tools to gather data "

"first. Afterwards I'll ask for a structured summary."

),

}

],

"followup": (

"Based on the tool results, produce the briefing as JSON matching the "

"country_report schema. Populate every required field and provide between "

"two and five highlights."

),

"validate": lambda parsed, tcs, raw: _validate_country_report(parsed, tcs),

}

def _validate_country_report(parsed, tcs):

names = [tc["function"]["name"] for tc in tcs]

for required_tool in ("get_country_stats", "get_climate_info"):

if required_tool not in names:

return False, f"missing tool call {required_tool!r}: got {names}"

required = {

"country", "capital", "population",

"climate_summary", "highlights", "suitable_for_tourism",

}

missing = required - parsed.keys()

if missing:

return False, f"missing report fields: {missing}"

if "norway" not in parsed["country"].lower():

return False, f"country should reference Norway: {parsed['country']!r}"

if "oslo" not in parsed["capital"].lower():

return False, f"capital should be Oslo: {parsed['capital']!r}"

if not isinstance(parsed["population"], int) or parsed["population"] < 1_000_000:

return False, f"population implausible: {parsed['population']!r}"

if not isinstance(parsed["climate_summary"], str) or not parsed["climate_summary"]:

return False, "climate_summary must be a non-empty string"

hls = parsed["highlights"]

if not isinstance(hls, list) or not (2 <= len(hls) <= 5):

return False, f"highlights length out of range: {hls!r}"

if not all(isinstance(h, str) and h for h in hls):

return False, "each highlight must be a non-empty string"

if not isinstance(parsed["suitable_for_tourism"], bool):

return False, f"suitable_for_tourism must be bool: {parsed['suitable_for_tourism']!r}"

return True, (

f"tools={names}; report for {parsed['country']} "

f"(pop {parsed['population']}, {len(hls)} highlights)"

)

# ---------------------------------------------------------------------------

# All test cases

# ---------------------------------------------------------------------------

ALL_TEST_CASES = [

BOOK_TEST_CASE,

SENTIMENT_TEST_CASE,

PRODUCT_JSON_OBJECT_TEST_CASE,

RECIPE_TEST_CASE,

SHOP_COMPARISON_TEST_CASE,

COUNTRY_REPORT_TEST_CASE,

]

# ---------------------------------------------------------------------------

# Entry point

# ---------------------------------------------------------------------------

def main():

parser = argparse.ArgumentParser(

description="Test llama-server structured-output capability."

)

parser.add_argument("--host", default="localhost")

parser.add_argument("--port", default=8080, type=int)

parser.add_argument(

"--no-stream", action="store_true", help="Disable streaming mode tests"

)

parser.add_argument(

"--stream-only", action="store_true", help="Only run streaming mode tests"

)