Skip to content

Commit 91ae770

Browse files
igordayenigordayen
authored
Move Jackson Converters from embabel-agent to embabel-common (#86)
1 parent b7679ae commit 91ae770

3 files changed

Lines changed: 246 additions & 0 deletions

File tree

‎embabel-common-ai/src/main/kotlin/com/embabel/common/ai/converters/FilteringJacksonOutputConverter.kt‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
/*
2+
* Copyright 2024-2025 Embabel Software, Inc.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.embabel.common.ai.converters
17+
18+
import com.fasterxml.jackson.databind.JsonNode
19+
import com.fasterxml.jackson.databind.ObjectMapper
20+
import com.fasterxml.jackson.databind.node.ObjectNode
21+
import org.springframework.core.ParameterizedTypeReference
22+
import java.lang.reflect.Type
23+
import java.util.function.Predicate
24+
25+
/**
26+
* Extension of [JacksonOutputConverter] that allows for filtering of properties of the generated object via a predicate.
27+
*/
28+
class FilteringJacksonOutputConverter<T> private constructor(
29+
type: Type,
30+
objectMapper: ObjectMapper,
31+
private val propertyFilter: Predicate<String>,
32+
) : JacksonOutputConverter<T>(type, objectMapper) {
33+
34+
constructor(
35+
clazz: Class<T>,
36+
objectMapper: ObjectMapper,
37+
propertyFilter: Predicate<String>,
38+
) : this(clazz as Type, objectMapper, propertyFilter)
39+
40+
constructor(
41+
typeReference: ParameterizedTypeReference<T>,
42+
objectMapper: ObjectMapper,
43+
propertyFilter: Predicate<String>,
44+
) : this(typeReference.type, objectMapper, propertyFilter)
45+
46+
override fun postProcessSchema(jsonNode: JsonNode) {
47+
val propertiesNode = jsonNode.get("properties") as? ObjectNode ?: return
48+
49+
val fieldNames = propertiesNode.fieldNames() as MutableIterator<String>
50+
while (fieldNames.hasNext()) {
51+
val fieldName = fieldNames.next()
52+
if (!this.propertyFilter.test(fieldName)) {
53+
fieldNames.remove()
54+
}
55+
}
56+
}
57+
}

‎embabel-common-ai/src/main/kotlin/com/embabel/common/ai/converters/JacksonOutputConverter.kt‎

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,123 @@
1+
/*
2+
* Copyright 2024-2025 Embabel Software, Inc.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.embabel.common.ai.converters
17+
18+
import com.fasterxml.jackson.core.JsonProcessingException
19+
import com.fasterxml.jackson.core.util.DefaultIndenter
20+
import com.fasterxml.jackson.core.util.DefaultPrettyPrinter
21+
import com.fasterxml.jackson.databind.JsonNode
22+
import com.fasterxml.jackson.databind.ObjectMapper
23+
import com.github.victools.jsonschema.generator.*
24+
import com.github.victools.jsonschema.module.jackson.JacksonModule
25+
import com.github.victools.jsonschema.module.jackson.JacksonOption
26+
import org.slf4j.Logger
27+
import org.slf4j.LoggerFactory
28+
import org.springframework.ai.converter.StructuredOutputConverter
29+
import org.springframework.ai.util.LoggingMarkers
30+
import org.springframework.core.ParameterizedTypeReference
31+
import java.lang.reflect.Type
32+
33+
/**
34+
* A Kotlin version of [org.springframework.ai.converter.BeanOutputConverter] that allows for customization
35+
* of the used schema via [postProcessSchema]
36+
*/
37+
open class JacksonOutputConverter<T> protected constructor(
38+
private val type: Type,
39+
val objectMapper: ObjectMapper,
40+
) : StructuredOutputConverter<T> {
41+
42+
constructor(
43+
clazz: Class<T>,
44+
objectMapper: ObjectMapper,
45+
) : this(clazz as Type, objectMapper)
46+
47+
constructor(
48+
typeReference: ParameterizedTypeReference<T>,
49+
objectMapper: ObjectMapper,
50+
) : this(typeReference.type, objectMapper)
51+
52+
protected val logger: Logger = LoggerFactory.getLogger(javaClass)
53+
54+
val jsonSchema: String by lazy {
55+
val jacksonModule = JacksonModule(
56+
JacksonOption.RESPECT_JSONPROPERTY_REQUIRED,
57+
JacksonOption.RESPECT_JSONPROPERTY_ORDER
58+
)
59+
val configBuilder = SchemaGeneratorConfigBuilder(
60+
SchemaVersion.DRAFT_2020_12,
61+
OptionPreset.PLAIN_JSON
62+
)
63+
.with(jacksonModule)
64+
.with(Option.FORBIDDEN_ADDITIONAL_PROPERTIES_BY_DEFAULT)
65+
val config = configBuilder.build()
66+
val generator = SchemaGenerator(config)
67+
val jsonNode: JsonNode = generator.generateSchema(this.type)
68+
postProcessSchema(jsonNode)
69+
val objectWriter = this.objectMapper.writer(
70+
DefaultPrettyPrinter()
71+
.withObjectIndenter(DefaultIndenter().withLinefeed(System.lineSeparator()))
72+
)
73+
try {
74+
objectWriter.writeValueAsString(jsonNode)
75+
} catch (e: JsonProcessingException) {
76+
logger.error("Could not pretty print json schema for jsonNode: {}", jsonNode)
77+
throw RuntimeException("Could not pretty print json schema for " + this.type, e)
78+
}
79+
}
80+
81+
/**
82+
* Empty template method that allows for customization of the JSON schema in subclasses.
83+
* @param jsonNode the JSON schema, in the form of a JSON node
84+
*/
85+
protected open fun postProcessSchema(jsonNode: JsonNode) {
86+
}
87+
88+
override fun convert(text: String): T? {
89+
val unwrapped = unwrapJson(text)
90+
try {
91+
return this.objectMapper.readValue<Any?>(unwrapped, this.objectMapper.constructType(this.type)) as T?
92+
} catch (e: JsonProcessingException) {
93+
logger.error(
94+
LoggingMarkers.SENSITIVE_DATA_MARKER,
95+
"Could not parse the given text to the desired target type: \"{}\" into {}", unwrapped, this.type
96+
)
97+
throw RuntimeException(e)
98+
}
99+
}
100+
101+
private fun unwrapJson(text: String): String {
102+
var result = text.trim()
103+
104+
if (result.startsWith("```") && result.endsWith("```")) {
105+
result = result.removePrefix("```json")
106+
.removePrefix("```")
107+
.removeSuffix("```")
108+
.trim()
109+
}
110+
111+
return result
112+
}
113+
114+
override fun getFormat(): String =
115+
"""|
116+
|Your response should be in JSON format.
117+
|Do not include any explanations, only provide a RFC8259 compliant JSON response following this format without deviation.
118+
|Do not include markdown code blocks in your response.
119+
|Remove the ```json markdown from the output.
120+
|Here is the JSON Schema instance your output must adhere to:
121+
|```${jsonSchema}```
122+
|""".trimMargin()
123+
}

‎embabel-common-ai/src/test/kotlin/com/embabel/common/ai/converters/FilteringJacksonOutputConverterTest.kt‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
/*
2+
* Copyright 2024-2025 Embabel Software, Inc.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.embabel.common.ai.converters
17+
18+
import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper
19+
import org.junit.jupiter.api.Assertions.assertFalse
20+
import org.junit.jupiter.api.Assertions.assertTrue
21+
import org.junit.jupiter.api.Test
22+
23+
class FilteringJacksonOutputConverterTest {
24+
25+
private val objectMapper = jacksonObjectMapper()
26+
27+
data class Person(
28+
val name: String,
29+
val age: Int,
30+
val email: String,
31+
val address: String
32+
)
33+
34+
@Test
35+
fun `test schema should include only specified properties`() {
36+
val converter = FilteringJacksonOutputConverter<Person>(
37+
clazz = Person::class.java,
38+
objectMapper = objectMapper,
39+
propertyFilter = { it == "name" || it == "age" }
40+
)
41+
42+
val schema = converter.jsonSchema
43+
44+
assertTrue(schema.contains("name"))
45+
assertTrue(schema.contains("age"))
46+
assertFalse(schema.contains("email"))
47+
assertFalse(schema.contains("address"))
48+
}
49+
50+
@Test
51+
fun `test schema should exclude specified properties`() {
52+
val converter = FilteringJacksonOutputConverter<Person>(
53+
clazz = Person::class.java,
54+
objectMapper = objectMapper,
55+
propertyFilter = { it != "email" && it != "address" }
56+
)
57+
58+
val schema = converter.jsonSchema
59+
60+
assertTrue(schema.contains("name"))
61+
assertTrue(schema.contains("age"))
62+
assertFalse(schema.contains("email"))
63+
assertFalse(schema.contains("address"))
64+
}
65+
66+
}

0 commit comments

Comments
 (0)