Customizing JSON Schema Generation for Callable Functions in Pydantic

[furlat]

Description:
We are working on a project that involves using Pydantic to generate JSON schemas for a game engine. The game engine includes various actions and interactions defined in the interactions.py file. Some of these actions and interactions involve callable functions that perform specific tasks or validations.

Problem:
We want to include detailed information about the callable functions in the generated JSON schema, such as the input arguments, input types, output type, docstring, and source code. This information will be helpful for the language model to understand the purpose and functionality of each callable function.

However, the default JSON schema generation in Pydantic does not include this level of detail for callable functions. We need a way to customize the JSON schema generation process to include the desired information.

Attempted Solution:
We tried to use the get_pydantic_json_schema method to customize the JSON schema generation for the callable functions. We created a decorator called json_schema_callable that wraps the callable functions and defines the get_pydantic_json_schema method.

Inside the get_pydantic_json_schema method, we used the inspect module to extract the input arguments, types, output type, docstring, and source code of the callable function. We then created a JsonSchemaValue object with this information and returned it as the custom JSON schema.

However, when we applied the @json_schema_callable decorator to the relevant callable functions in the interactions.py file and tried to generate the JSON schema using MoveStep.model_json_schema(), we encountered errors related to the core_schema.CallableSchema and the JSON schema generation process.

Possible Solutions:
After reviewing the Pydantic documentation, we found a few potential solutions to customize the JSON schema generation for callable functions:

Implementing get_pydantic_core_schema:
We can implement the get_pydantic_core_schema method on the callable functions or classes to customize the core schema generation.
This method receives the source_type and handler as arguments and returns a CoreSchema instance.
We can use this method to define custom validation logic and modify the core schema based on the callable function's properties.
Reference: Implementing get_pydantic_core_schema
Using the WithJsonSchema annotation:
We can use the WithJsonSchema annotation to override the generated JSON schema for a given type without the need to implement get_pydantic_core_schema or get_pydantic_json_schema on the type itself.
This annotation allows us to set a JSON schema for types that would otherwise raise errors when producing a JSON schema, such as Callable.
We can define a custom JSON schema for the callable functions using the WithJsonSchema annotation.
Reference: WithJsonSchema annotation
Customizing the JSON Schema Generation Process:
We can create a custom subclass of GenerateJsonSchema and override the relevant methods to modify the JSON schema generation process.
This approach allows us to customize the generation of JSON schemas for specific types or fields.
We can define custom logic to handle callable functions and include the desired information in the generated JSON schema.
Reference: Customizing the JSON Schema Generation Process
Next Steps:

Investigate the get_pydantic_core_schema method and explore how it can be used to customize the core schema generation for callable functions.
Experiment with the WithJsonSchema annotation to define custom JSON schemas for the callable functions.
Consider creating a custom subclass of GenerateJsonSchema to modify the JSON schema generation process specifically for callable functions.
Test the different approaches and evaluate which one provides the most suitable solution for including the desired information in the generated JSON schema.
Update the interactions.py file and the json_schema_callable decorator based on the chosen approach.
Verify that the generated JSON schema includes the input arguments, types, output type, docstring, and source code of the callable functions.
By exploring these potential solutions from the Pydantic documentation, we can find a way to customize the JSON schema generation process to include the necessary information for the callable functions in our game engine.

Add to Conversation