Add support for structured outputs

[pminervini]

As in https://docs.vllm.ai/en/latest/features/structured_outputs/#online-serving-openai-api and https://developers.openai.com/api/docs/guides/structured-outputs

I tried it with ds4 using the latest code from main but it doesn't seem to be working fine (both DGX Spark and m3 max) -- here's the test:

#!/usr/bin/env python3
from openai import OpenAI
from pydantic import BaseModel, ConfigDict


BASE_URL = "http://127.0.0.1:8000/v1"
API_KEY = "not-needed"
MODEL = "default"


class CalendarEvent(BaseModel):
    model_config = ConfigDict(extra="forbid")

    name: str
    date: str
    participants: list[str]


def main() -> None:
    client = OpenAI(base_url=BASE_URL, api_key=API_KEY)

    print(f"[so-test] base_url={BASE_URL}")
    print(f"[so-test] model={MODEL}")
    print("[so-test] calling client.responses.parse(..., text_format=CalendarEvent)")

    response = client.responses.parse(
        model=MODEL,
        input=[
            {"role": "system", "content": "Extract the event information."},
            {
                "role": "user",
                "content": "Alice and Bob are going to a science fair on Friday.",
            },
        ],
        text_format=CalendarEvent,
    )

    event = response.output_parsed
    print("[so-test] parsed object:")
    print(event.model_dump_json(indent=2))
    print("[so-test] raw response id:", getattr(response, "id", None))


if __name__ == "__main__":
    main()

Output:

$ python3 so-test-cli.py 
[so-test] base_url=http://127.0.0.1:8000/v1
[so-test] model=default
[so-test] calling client.responses.parse(..., text_format=CalendarEvent)
Traceback (most recent call last):
[..]
pydantic_core._pydantic_core.ValidationError: 6 validation errors for CalendarEvent
people
  Extra inputs are not permitted [type=extra_forbidden, input_value=['Alice', 'Bob'], input_type=list]
    For further information visit https://errors.pydantic.dev/2.13/v/extra_forbidden
event
  Extra inputs are not permitted [type=extra_forbidden, input_value='science fair', input_type=str]
    For further information visit https://errors.pydantic.dev/2.13/v/extra_forbidden
day
  Extra inputs are not permitted [type=extra_forbidden, input_value='Friday', input_type=str]
    For further information visit https://errors.pydantic.dev/2.13/v/extra_forbidden
name
  Field required [type=missing, input_value={'people': ['Alice', 'Bob... fair', 'day': 'Friday'}, input_type=dict]
    For further information visit https://errors.pydantic.dev/2.13/v/missing
date
  Field required [type=missing, input_value={'people': ['Alice', 'Bob... fair', 'day': 'Friday'}, input_type=dict]
    For further information visit https://errors.pydantic.dev/2.13/v/missing
participants
  Field required [type=missing, input_value={'people': ['Alice', 'Bob... fair', 'day': 'Friday'}, input_type=dict]
    For further information visit https://errors.pydantic.dev/2.13/v/missing

A bit swamped but happy to look into this at some point

[pminervini]

Not sure how robust this is but it seems to work for simple schemas (then one can think of implementing a retry logic on top of it but I wouldn't trust that with my life):

#!/usr/bin/env python3
import json

import openai
from pydantic import BaseModel, ConfigDict


BASE_URL = "http://127.0.0.1:8000/v1"
API_KEY = "not-needed"
MODEL = "default"


class CalendarEvent(BaseModel):
    model_config = ConfigDict(extra="forbid")

    name: str
    date: str
    participants: list[str]


def main() -> None:
    client = openai.OpenAI(base_url=BASE_URL, api_key=API_KEY)
    schema = CalendarEvent.model_json_schema()

    print(f"[so-test-json] base_url={BASE_URL}")
    print(f"[so-test-json] model={MODEL}")
    print('[so-test-json] calling responses.create with text={"format": {"type": "json_object"}}')

    response = client.responses.create(
        model=MODEL,
        instructions=(
            "Extract the event information and return one JSON object "
            "that matches this JSON schema. Return only the JSON object, "
            "not the schema and not markdown.\n\n"
            f"{json.dumps(schema, indent=2)}"
        ),
        input="Alice and Bob are going to a science fair on Friday.",
        text={"format": {"type": "json_object"}},
        temperature=0.0,
        max_output_tokens=500,
    )

    raw = response.output_text
    event = CalendarEvent.model_validate_json(raw)

    print("[so-test-json] parsed object:")
    print(event.model_dump_json(indent=2))
    print("[so-test-json] raw response id:", getattr(response, "id", None))


if __name__ == "__main__":
    main()

% python3 so-test-json-cli.py
[so-test-json] base_url=http://127.0.0.1:8000/v1
[so-test-json] model=default
[so-test-json] calling responses.create with text={"format": {"type": "json_object"}}
[so-test-json] parsed object:
{
  "name": "science fair",
  "date": "Friday",
  "participants": [
    "Alice",
    "Bob"
  ]
}
[so-test-json] raw response id: resp_46b085d2c52dd83d3158e9bb

[pminervini]

Alternate approach:

#!/usr/bin/env python3
import openai
from pydantic import BaseModel, ConfigDict


BASE_URL = "http://127.0.0.1:8000/v1"
API_KEY = "not-needed"
MODEL = "default"
TOOL_NAME = "record_calendar_event"


class CalendarEvent(BaseModel):
    model_config = ConfigDict(extra="forbid")

    name: str
    date: str
    participants: list[str]


def main() -> None:
    client = openai.OpenAI(base_url=BASE_URL, api_key=API_KEY)
    tool = {
        "type": "function",
        "name": TOOL_NAME,
        "description": "Record the extracted calendar event.",
        "parameters": CalendarEvent.model_json_schema(),
        "strict": True,
    }

    print(f"[so-test-ts] base_url={BASE_URL}")
    print(f"[so-test-ts] model={MODEL}")
    print(f"[so-test-ts] calling responses.create with strict tool={TOOL_NAME}")

    response = client.responses.create(
        model=MODEL,
        instructions=(
            f"You must call the {TOOL_NAME} function. Do not answer with text. "
            "Extract only facts from the user message."
        ),
        input="Alice and Bob are going to a science fair on Friday.",
        tools=[tool],
        tool_choice="auto",
        parallel_tool_calls=False,
        temperature=0.0,
        max_output_tokens=500,
    )

    call = next(item for item in response.output if item.type == "function_call")
    print(f"[so-test-ts] tool={call.name}")

    event = CalendarEvent.model_validate_json(call.arguments)
    print("[so-test-ts] parsed object:")
    print(event.model_dump_json(indent=2))
    print("[so-test-ts] raw response id:", getattr(response, "id", None))


if __name__ == "__main__":
    main()

% python3 so-test-ts-cli.py
[so-test-ts] base_url=http://127.0.0.1:8000/v1
[so-test-ts] model=default
[so-test-ts] calling responses.create with strict tool=record_calendar_event
[so-test-ts] tool=record_calendar_event
[so-test-ts] parsed object:
{
  "name": "science fair",
  "date": "Friday",
  "participants": [
    "Alice",
    "Bob"
  ]
}
[so-test-ts] raw response id: resp_014ae3e48659504136f6e2dc

[pminervini]

Here's a tentative fix that uses constrained decoding similarly to llama.cpp (since just prompting the LLM with a JSON schema and hoping for the best may not be robust enough): https://github.com/pminervini/ds4/tree/issue-210-structured-json

[fry69]

I know the "reinventing the wheel" karma is strong in this project (for good, valid reasons), but for structured output it may make sense to use a project that mostly solved this problem and get used by almost all other LLM engines ("xgrammar")?

It is huge, I admit, but it seems to come with C bindings? -> https://github.com/mlc-ai/xgrammar/tree/main/include/xgrammar

The upside is that this enables all kinds of structured output, not just JSON.

[pminervini]

@fry69 yep this PR is a bit bloated, xgrammar would do the trick but not sure how @antirez feels about introducing deps -- if that's ok I can try to give it a stab

[fry69]

xgrammar would do the trick but not sure how

I come from the perspective of how -> https://github.com/jundot/omlx implements this, see here -> https://github.com/search?q=repo%3Ajundot%2Fomlx%20xgrammar&type=code

It is not even enabled by default in oMLX because xgrammar is so huge, but it might help for inspiration how this can be done?

Since structured outputs tend to be niche (but very useful!) and they will come with overhead, no matter what, it may make sense to go the same route -> optional, maybe even an outside extension?

[pminervini]

Yep, I was also checking how llama.cpp does it -- they don't use xgrammar, but they let you redirect requests to https://github.com/guidance-ai/llguidance via -DLLAMA_LLGUIDANCE=ON

Shall we try something similar? I'm not a fan of this PR because it feels very bloated; switching to xgrammar may help maintain the codebase lean

[fry69]

Looks like -> https://github.com/guidance-ai/llguidance is the better choice in this case, more native and faster, and with llama.cpp a widely used implementation using this library.

I hadn't researched the available options. My main motivation is pointing at an existing solution, since getting structured outputs right is finicky and the time may be better spent on other problems.

But of course in this brave new AI world reimplementing for the sake of reimplementing maybe worth it anyway? E.g. adopting a library of tests and go for it?