Skip to content

Commit

Permalink
added SchemaConverter interface for controlling Schema generation fro…
Browse files Browse the repository at this point in the history
…m classes
  • Loading branch information
asaf-wizzdi committed Jul 18, 2024
1 parent 5db4c2c commit 13a9a3a
Show file tree
Hide file tree
Showing 7 changed files with 181 additions and 31 deletions.
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
package io.github.sashirestela.openai.common.function;

import io.github.sashirestela.openai.support.JsonSchemaUtil;
import lombok.Builder;
import lombok.Getter;
import lombok.NonNull;
Expand All @@ -15,5 +16,7 @@ public class FunctionDef {

@NonNull
private Class<? extends Functional> functionalClass;
@Builder.Default
private SchemaConverter schemaConverter= JsonSchemaUtil.defaultConverter;

}
Original file line number Diff line number Diff line change
@@ -0,0 +1,10 @@
package io.github.sashirestela.openai.common.function;

import com.fasterxml.jackson.databind.JsonNode;

public interface SchemaConverter {

JsonNode convert(Class<?> c);


}
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ public static Tool function(FunctionDef function) {
new ToolFunctionDef(
function.getName(),
function.getDescription(),
JsonSchemaUtil.classToJsonSchema(function.getFunctionalClass())));
function.getSchemaConverter().convert(function.getFunctionalClass())));
}

@AllArgsConstructor
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,43 @@
package io.github.sashirestela.openai.support;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.github.victools.jsonschema.generator.*;
import com.github.victools.jsonschema.module.jackson.JacksonModule;
import com.github.victools.jsonschema.module.jackson.JacksonOption;
import io.github.sashirestela.openai.SimpleUncheckedException;
import io.github.sashirestela.openai.common.function.SchemaConverter;

import static io.github.sashirestela.openai.support.JsonSchemaUtil.JSON_EMPTY_CLASS;

public class DefaultSchemaConverter implements SchemaConverter {
private final SchemaGenerator schemaGenerator;
private final ObjectMapper objectMapper;

public DefaultSchemaConverter() {
objectMapper = new ObjectMapper();
var jacksonModule = new JacksonModule(JacksonOption.RESPECT_JSONPROPERTY_REQUIRED,
JacksonOption.RESPECT_JSONPROPERTY_ORDER);
var configBuilder = new SchemaGeneratorConfigBuilder(SchemaVersion.DRAFT_2020_12,
OptionPreset.PLAIN_JSON)
.with(jacksonModule)
.without(Option.SCHEMA_VERSION_INDICATOR);
var config = configBuilder.build();
schemaGenerator = new SchemaGenerator(config);
}

@Override
public JsonNode convert(Class<?> clazz) {
JsonNode jsonSchema;
try {
jsonSchema = schemaGenerator.generateSchema(clazz);
if (jsonSchema.get("properties") == null) {
jsonSchema = objectMapper.readTree(JSON_EMPTY_CLASS);
}

} catch (Exception e) {
throw new SimpleUncheckedException("Cannot generate the Json Schema for the class {0}.", clazz.getName(), e);
}
return jsonSchema;
}
}
Original file line number Diff line number Diff line change
@@ -1,45 +1,19 @@
package io.github.sashirestela.openai.support;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.github.victools.jsonschema.generator.Option;
import com.github.victools.jsonschema.generator.OptionPreset;
import com.github.victools.jsonschema.generator.SchemaGenerator;
import com.github.victools.jsonschema.generator.SchemaGeneratorConfigBuilder;
import com.github.victools.jsonschema.generator.SchemaVersion;
import com.github.victools.jsonschema.module.jackson.JacksonModule;
import com.github.victools.jsonschema.module.jackson.JacksonOption;
import io.github.sashirestela.openai.SimpleUncheckedException;
import io.github.sashirestela.openai.common.function.SchemaConverter;

public class JsonSchemaUtil {

public static final SchemaConverter defaultConverter= new DefaultSchemaConverter();

public static final String JSON_EMPTY_CLASS = "{\"type\":\"object\",\"properties\":{}}";
private static ObjectMapper objectMapper = new ObjectMapper();

private JsonSchemaUtil() {
}

public static JsonNode classToJsonSchema(Class<?> clazz) {
JsonNode jsonSchema = null;
try {
var jacksonModule = new JacksonModule(JacksonOption.RESPECT_JSONPROPERTY_REQUIRED,
JacksonOption.RESPECT_JSONPROPERTY_ORDER);
var configBuilder = new SchemaGeneratorConfigBuilder(SchemaVersion.DRAFT_2020_12,
OptionPreset.PLAIN_JSON)
.with(jacksonModule)
.without(Option.SCHEMA_VERSION_INDICATOR);
var config = configBuilder.build();
var generator = new SchemaGenerator(config);
jsonSchema = generator.generateSchema(clazz);
if (jsonSchema.get("properties") == null) {
jsonSchema = objectMapper.readTree(JSON_EMPTY_CLASS);
}

} catch (Exception e) {
throw new SimpleUncheckedException("Cannot generate the Json Schema for the class {0}.",
clazz.getName(), e);
}
return jsonSchema;
return defaultConverter.convert(clazz);
}

}
Original file line number Diff line number Diff line change
@@ -0,0 +1,46 @@
package io.github.sashirestela.openai.support;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.github.victools.jsonschema.generator.*;
import com.github.victools.jsonschema.generator.Module;
import com.github.victools.jsonschema.module.jackson.JacksonModule;
import com.github.victools.jsonschema.module.jackson.JacksonOption;
import io.github.sashirestela.openai.SimpleUncheckedException;
import io.github.sashirestela.openai.common.function.SchemaConverter;

public class CustomSchemaConverter implements SchemaConverter {
private final SchemaGenerator schemaGenerator;
private final ObjectMapper objectMapper;
public static final String JSON_EMPTY_CLASS = "{\"type\":\"object\",\"properties\":{}}";

public CustomSchemaConverter() {
objectMapper = new ObjectMapper();
var jacksonModule = new JacksonModule(JacksonOption.RESPECT_JSONPROPERTY_REQUIRED,
JacksonOption.RESPECT_JSONPROPERTY_ORDER);
var configBuilder = new SchemaGeneratorConfigBuilder(SchemaVersion.DRAFT_2020_12,
OptionPreset.PLAIN_JSON)
.with(jacksonModule)
.with(builder -> builder.forTypesInGeneral().withTypeAttributeOverride(
(collectedTypeAttributes, scope, context) -> collectedTypeAttributes.put("myCustomProperty",true)))
.without(Option.SCHEMA_VERSION_INDICATOR);
var config = configBuilder.build();
schemaGenerator = new SchemaGenerator(config);
}

@Override
public JsonNode convert(Class<?> clazz) {
JsonNode jsonSchema;
try {
jsonSchema = schemaGenerator.generateSchema(clazz);
if (jsonSchema.get("properties") == null) {
jsonSchema = objectMapper.readTree(JSON_EMPTY_CLASS);
}

} catch (Exception e) {
throw new SimpleUncheckedException("Cannot generate the Json Schema for the class {0}.", clazz.getName(), e);
}
return jsonSchema;
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,74 @@
package io.github.sashirestela.openai.support;

import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.annotation.JsonPropertyOrder;
import com.fasterxml.jackson.databind.JsonNode;
import io.github.sashirestela.openai.common.function.SchemaConverter;
import lombok.AllArgsConstructor;
import lombok.Getter;
import lombok.NoArgsConstructor;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;

import static io.github.sashirestela.openai.support.JsonSchemaUtil.JSON_EMPTY_CLASS;
import static org.junit.jupiter.api.Assertions.assertEquals;

class CustomSchemaConverterTest {

private static SchemaConverter schemaConverter=new CustomSchemaConverter();


@Test
void shouldGenerateFullJsonSchemaWhenClassHasSomeFields() {
var actualJsonSchema = schemaConverter.convert(TestClass.class).toString();
var expectedJsonSchema = "{\"type\":\"object\",\"properties\":{\"first\":{\"type\":\"string\",\"myCustomProperty\":true},\"second\":{\"type\":\"integer\",\"myCustomProperty\":true}},\"required\":[\"first\"],\"myCustomProperty\":true}";
assertEquals(expectedJsonSchema, actualJsonSchema);
}

@Test
void shouldGenerateEmptyJsonSchemaWhenClassHasNoFields() {
var actualJsonSchema = schemaConverter.convert(EmptyClass.class).toString();
var expectedJsonSchema = JSON_EMPTY_CLASS;
assertEquals(expectedJsonSchema, actualJsonSchema);
}

@Test
void shouldGenerateOrderedJsonSchemaWhenClassHasJsonPropertyOrderAnnotation() {
var actualJsonSchema = schemaConverter.convert(OrderedTestClass.class).toString();
var expectedJsonSchema = "{\"type\":\"object\",\"properties\":{\"first\":{\"type\":\"string\",\"myCustomProperty\":true},\"second\":{\"type\":\"integer\",\"myCustomProperty\":true},\"third\":{\"type\":\"string\",\"myCustomProperty\":true}},\"required\":[\"first\"],\"myCustomProperty\":true}";
assertEquals(expectedJsonSchema, actualJsonSchema);
}

@NoArgsConstructor
@AllArgsConstructor
@Getter
static class TestClass {

@JsonProperty(required = true)
public String first;

public Integer second;

}

static class EmptyClass {
}

@NoArgsConstructor
@AllArgsConstructor
@Getter
@JsonPropertyOrder({ "first", "second", "third" })
static class OrderedTestClass {

@JsonProperty(required = true)
public String first;

public Integer second;

public String third;

}



}

0 comments on commit 13a9a3a

Please sign in to comment.