[FEAT] Support JSON Schema in Responses

[riedgar-ms]

Description

The OpenAI Responses API provides for structured output with JSON schema. This change:

• Fixes support in OpenAIResponsesTarget when JSON output is requested without a schema
• Allows the user to supply a schema via prompt_metadata
• Adds a test for each scenario

Tests and Documentation

I have added a test for both scenarios (i.e. with and without schema), and updated the appropriate notebooks.

[romanlutz]

Do you expect a different schema per message sent, or just one schema for one target? I think that's the main decision point. If it's the latter then we can pass it to the target. If it's the former that makes it trickier and I would probably suggest something like what you're doing here.

[riedgar-ms]

Schema is per-response. You might want different JSONs as a conversation develops, and you also probably don't want a JSON every time. This came out of finding that the SelfAskScorer wasn't working with an OpenAIResponsesTarget; I've not tracked down that prompt yet (ISTR YAML files), but for all of them, anything which asks for a JSON response could have the required schema stored there too, and passed along as needed.

I was wondering if I could store the schema as a Python object (which would be serialisable to JSON), but I've not dug through things enough to know if prompt_metadata has to have str or int values.

[romanlutz]

Re: Response target not working with the scorer

We don't currently enforce a json format, we just use the other option (I don't recall the name offhand, but switching to schema was definitely on our list of things to do). What exactly is failing currently?

Maybe this should be a new parameter on a message piece... curious what @rlundeen2 thinks.

[riedgar-ms]

The specific bug is that the Responses API doesn't use response_format in the body, but the alternative I show. There are also two options: json_object which doesn't enforce a schema (although may enforce JSON output; I haven't dug enough through the code to be sure), and json_schema where you can pass an actual JSON schema object.

[romanlutz]

The specific bug is that the Responses API doesn't use response_format in the body, but the alternative I show. There are also two options: json_object which doesn't enforce a schema (although may enforce JSON output; I haven't dug enough through the code to be sure), and json_schema where you can pass an actual JSON schema object.

Yes we support json_object although it's more of a suggestion than enforced output format. We noticed json_schema a while ago but the problem was always how to pass the schema. I'll ping @rlundeen2 and @bashirpartovi internally to see if anyone has thoughts.

[riedgar-ms]

@romanlutz @rlundeen2 , I've just pushed changes which modify is_response_format_json() so that if a JSON schema is supplied, then it will return that schema. Now I'm not a huge fan of functions having different return types like this, but changing the return type to a tuple didn't feel much better.

I think I've caught all the places this could have affected.....

[bashirpartovi]

Thanks @riedgar-ms for helping fix this issue. However, I do not agree with this approach because it violates least surprise principle and does not provide a clean design. I believe all of this can be encapsulated much cleaner (this is just to demonstrate an approach):

We could define a class to encapsulate json response configurations such as follows:

from __future__ import annotations
...

@dataclass
class JsonResponseConfig:

    enabled: bool = False
    schema: Optional[Dict[str, Any]] = None
    schema_name: str = "CustomSchema"
    strict: bool = True
    
    @classmethod
    def from_metadata(cls, metadata: Optional[Dict[str, Any]]) -> "JsonResponseConfig":
        if not metadata:
            return cls(enabled=False)
            
        response_format = metadata.get("response_format")
        if response_format != "json":
            return cls(enabled=False)
            
        schema_val = metadata.get("json_schema")
        if schema_val:
            if isinstance(schema_val, str):
                try:
                    schema = json.loads(schema_val) if schema_val else None
                except json.JSONDecodeError:
                    raise ValueError(f"Invalid JSON schema provided: {schema_val}")
            else:
                schema = schema_val
                
            return cls(
                enabled=True,
                schema=schema,
                schema_name=metadata.get("schema_name", "CustomSchema"),
                strict=metadata.get("strict", True)
            )
        
        return cls(enabled=True)

We also need to add a new method to the prompt chat target to handle retrieving the json config:

class PromptChatTarget(PromptTarget):
    ...

    def is_response_format_json(self, *, message_piece: MessagePiece) -> bool:
        config = self.get_json_response_config(message_piece=message_piece)
        return config.enabled

    def get_json_response_config(self, *, message_piece: MessagePiece) -> JsonResponseConfig:
        config = JsonResponseConfig.from_metadata(message_piece.prompt_metadata)
        
        if config.enabled and not self.is_json_response_supported():
            target_name = self.get_identifier()["__type__"]
            raise ValueError(f"This target {target_name} does not support JSON response format.")
            
        return config

We also need to change the base target for open ai to update the signature with the json config class:

class OpenAIChatTargetBase(OpenAITarget, PromptChatTarget):
    ...

    async def _construct_request_body(
        self,
        *,
        conversation: MutableSequence[Message],
        json_config: JsonResponseConfig
    ) -> dict:
        raise NotImplementedError

OpenAI chat target becomes:

class OpenAIChatTarget(OpenAIChatTargetBase):
    ...
    async def _construct_request_body(
        self,
        *,
        conversation: MutableSequence[Message],
        json_config: JsonResponseConfig
    ) -> dict:
        messages = await self._build_chat_messages_async(conversation)
        response_format = self._build_response_format(json_config)
        
        body_parameters = {
            "model": self._model_name,
            "max_completion_tokens": self._max_completion_tokens,
            "temperature": self._temperature,
            "top_p": self._top_p,
            "frequency_penalty": self._frequency_penalty,
            "presence_penalty": self._presence_penalty,
            "logit_bias": self._logit_bias,
            "stream": False,
            "seed": self._seed,
            "n": self._n,
            "messages": messages,
            "response_format": response_format,
        }

        if self._extra_body_parameters:
            body_parameters.update(self._extra_body_parameters)

        return {k: v for k, v in body_parameters.items() if v is not None}
    
    def _build_response_format(self, json_config: JsonResponseConfig) -> Optional[Dict[str, Any]]:
        if not json_config.enabled:
            return None
            
        if json_config.schema:
            return {
                "type": "json_schema",
                "json_schema": {
                    "name": json_config.schema_name,
                    "schema": json_config.schema,
                    "strict": json_config.strict
                }
            }
        
        return {"type": "json_object"}

Open AI response target:

class OpenAIResponseTarget(OpenAIChatTargetBase):
    ...

    async def _construct_request_body(
        self,
        *,
        conversation: MutableSequence[Message],
        json_config: JsonResponseConfig
    ) -> dict:
        input_items = await self._build_input_for_multi_modal_async(conversation)

        text_format = self._build_text_format(json_config)
        
        body_parameters = {
            "model": self._model_name,
            "max_output_tokens": self._max_output_tokens,
            "temperature": self._temperature,
            "top_p": self._top_p,
            "stream": False,
            "input": input_items,
            "text": text_format,
        }

        if self._extra_body_parameters:
            body_parameters.update(self._extra_body_parameters)

        return {k: v for k, v in body_parameters.items() if v is not None}
    
    def _build_text_format(self, json_config: JsonResponseConfig) -> Optional[Dict[str, Any]]:
        if not json_config.enabled:
            return None
            
        if json_config.schema:
            return {
                "format": {
                    "type": "json_schema",
                    "json_schema": {
                        "name": json_config.schema_name,
                        "schema": json_config.schema,
                        "strict": json_config.strict
                    }
                }
            }
        
        logger.info("Using json_object format without schema - consider providing a schema for better results")
        return {"format": {"type": "json_object"}}

OpenAI chat target base:

class OpenAIChatTargetBase(OpenAITarget, PromptChatTarget):
   ...
    
    async def send_prompt_async(self, *, message: Message) -> Message:
        self._validate_request(message=message)
        conversation = self._get_conversation(message=message)
        
        json_config = JsonResponseConfig(enabled=False)
        if message.message_pieces:
            last_piece = message.message_pieces[-1]
            json_config = self.get_json_response_config(message_piece=last_piece)
        
        request_body = await self._construct_request_body(
            conversation=conversation,
            json_config=json_config
        )
        .... 
        # rest of the method

[riedgar-ms]

@bashirpartovi happy to make changes along those lines

[rlundeen2]

Agreed with this approach Bashir suggested. I know it's a bigger change but I think will be a lot more maintainable.