Skip to content

Commit d90c8e7

Browse files
jgindinjgindin
committed
Treeshake common types schema
Port of Python SDK commit 1ea689d
1 parent 12e12ca commit d90c8e7

3 files changed

Lines changed: 189 additions & 8 deletions

File tree

‎agent_sdks/kotlin/src/main/kotlin/com/google/a2ui/core/schema/Catalog.kt‎

Lines changed: 68 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,9 @@ data class A2uiCatalog(
5656
@JvmField val catalogSchema: JsonObject,
5757
) {
5858

59-
private val logger = Logger.getLogger(A2uiCatalog::class.java.name)
59+
companion object {
60+
private val logger = Logger.getLogger(A2uiCatalog::class.java.name)
61+
}
6062

6163
val validator: A2uiValidator by lazy { A2uiValidator(this) }
6264

@@ -76,7 +78,7 @@ data class A2uiCatalog(
7678
* @return A copy of the catalog with only allowed components.
7779
*/
7880
fun withPrunedComponents(allowedComponents: List<String>): A2uiCatalog {
79-
if (allowedComponents.isEmpty()) return this
81+
if (allowedComponents.isEmpty()) return this.withPrunedCommonTypes()
8082

8183
val schemaCopy = catalogSchema.toMutableMap()
8284

@@ -95,7 +97,68 @@ data class A2uiCatalog(
9597
}
9698
}
9799

98-
return copy(catalogSchema = JsonObject(schemaCopy))
100+
return copy(catalogSchema = JsonObject(schemaCopy)).withPrunedCommonTypes()
101+
}
102+
103+
/** Returns a new catalog with unused common types pruned from the schema. */
104+
private fun withPrunedCommonTypes(): A2uiCatalog {
105+
val defs = commonTypesSchema["\$defs"] as? JsonObject ?: return this
106+
if (defs.isEmpty()) return this
107+
108+
fun collectRefs(element: JsonElement, refs: MutableSet<String>) {
109+
when (element) {
110+
is JsonObject -> {
111+
for ((k, v) in element) {
112+
if (k == "\$ref" && v is JsonPrimitive && v.isString) {
113+
refs.add(v.content)
114+
} else {
115+
collectRefs(v, refs)
116+
}
117+
}
118+
}
119+
is JsonArray -> {
120+
for (item in element) {
121+
collectRefs(item, refs)
122+
}
123+
}
124+
else -> {}
125+
}
126+
}
127+
128+
val visitedDefs = mutableSetOf<String>()
129+
val queue = ArrayDeque<String>()
130+
131+
val externalRefs = mutableSetOf<String>()
132+
collectRefs(catalogSchema, externalRefs)
133+
collectRefs(serverToClientSchema, externalRefs)
134+
135+
val prefix = "common_types.json#/\$defs/"
136+
for (ref in externalRefs) {
137+
if (ref.startsWith(prefix)) {
138+
queue.add(ref.substring(prefix.length))
139+
}
140+
}
141+
142+
while (queue.isNotEmpty()) {
143+
val defName = queue.removeFirst()
144+
if (defs.containsKey(defName) && visitedDefs.add(defName)) {
145+
val defElement = defs[defName]!!
146+
val internalRefs = mutableSetOf<String>()
147+
collectRefs(defElement, internalRefs)
148+
val internalPrefix = "#/\$defs/"
149+
for (ref in internalRefs) {
150+
if (ref.startsWith(internalPrefix)) {
151+
queue.add(ref.substring(internalPrefix.length))
152+
}
153+
}
154+
}
155+
}
156+
157+
val newDefs = JsonObject(defs.filterKeys { it in visitedDefs })
158+
val newCommonTypes =
159+
JsonObject(commonTypesSchema.toMutableMap().apply { put("\$defs", newDefs) })
160+
161+
return copy(commonTypesSchema = newCommonTypes)
99162
}
100163

101164
private fun pruneAnyComponentOneOf(
@@ -127,10 +190,8 @@ data class A2uiCatalog(
127190
appendLine("### Server To Client Schema:")
128191
appendLine(jsonFmt.encodeToString(JsonElement.serializer(), serverToClientSchema))
129192

130-
if (
131-
commonTypesSchema.isNotEmpty() &&
132-
(commonTypesSchema["\$defs"] as? JsonObject)?.isEmpty() != true
133-
) {
193+
val defs = commonTypesSchema["\$defs"] as? JsonObject
194+
if (!defs.isNullOrEmpty()) {
134195
appendLine("\n### Common Types Schema:")
135196
appendLine(jsonFmt.encodeToString(JsonElement.serializer(), commonTypesSchema))
136197
}

‎agent_sdks/kotlin/src/test/kotlin/com/google/a2ui/conformance/ConformanceTest.kt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -442,7 +442,7 @@ class ConformanceTest {
442442

443443
val dummyCatalog =
444444
Json.parseToJsonElement(
445-
"""{"catalogId": "https://a2ui.org/specification/v0_8/standard_catalog_definition.json", "components": {"Text": {}}}"""
445+
"{\"catalogId\": \"https://a2ui.org/specification/v0_8/standard_catalog_definition.json\", \"components\": {\"Text\": {\"\$ref\": \"common_types.json#/\$defs/DynamicString\"}}}"
446446
)
447447
.jsonObject
448448
val dummyConfig =

‎agent_sdks/kotlin/src/test/kotlin/com/google/a2ui/core/schema/CatalogPruningTest.kt‎

Lines changed: 120 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,120 @@
1+
/*
2+
* Copyright 2026 Google LLC
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+
* https://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+
17+
package com.google.a2ui.core.schema
18+
19+
import kotlin.test.Test
20+
import kotlin.test.assertFalse
21+
import kotlin.test.assertTrue
22+
import kotlinx.serialization.json.JsonObject
23+
import kotlinx.serialization.json.JsonPrimitive
24+
import kotlinx.serialization.json.jsonObject
25+
26+
class CatalogPruningTest {
27+
28+
@Test
29+
fun omitsEmptyCommonTypesFromLlmInstructions() {
30+
val catalogEmpty =
31+
A2uiCatalog(
32+
version = A2uiVersion.VERSION_0_9,
33+
name = "test",
34+
serverToClientSchema = JsonObject(emptyMap()),
35+
commonTypesSchema = JsonObject(emptyMap()),
36+
catalogSchema = JsonObject(mapOf(A2uiConstants.CATALOG_ID_KEY to JsonPrimitive("test_id"))),
37+
)
38+
assertFalse(catalogEmpty.renderAsLlmInstructions().contains("### Common Types Schema:"))
39+
40+
val catalogNoDefs =
41+
A2uiCatalog(
42+
version = A2uiVersion.VERSION_0_9,
43+
name = "test",
44+
serverToClientSchema = JsonObject(emptyMap()),
45+
commonTypesSchema = JsonObject(mapOf("something" to JsonPrimitive("else"))),
46+
catalogSchema = JsonObject(mapOf(A2uiConstants.CATALOG_ID_KEY to JsonPrimitive("test_id"))),
47+
)
48+
assertFalse(catalogNoDefs.renderAsLlmInstructions().contains("### Common Types Schema:"))
49+
50+
val catalogEmptyDefs =
51+
A2uiCatalog(
52+
version = A2uiVersion.VERSION_0_9,
53+
name = "test",
54+
serverToClientSchema = JsonObject(emptyMap()),
55+
commonTypesSchema = JsonObject(mapOf("\$defs" to JsonObject(emptyMap()))),
56+
catalogSchema = JsonObject(mapOf(A2uiConstants.CATALOG_ID_KEY to JsonPrimitive("test_id"))),
57+
)
58+
assertFalse(catalogEmptyDefs.renderAsLlmInstructions().contains("### Common Types Schema:"))
59+
}
60+
61+
@Test
62+
fun prunesUnusedCommonTypesCorrectly() {
63+
val commonTypes =
64+
JsonObject(
65+
mapOf(
66+
"\$defs" to
67+
JsonObject(
68+
mapOf(
69+
"TypeForCompA" to
70+
JsonObject(
71+
mapOf(
72+
"type" to JsonPrimitive("string"),
73+
"\$ref" to JsonPrimitive("#/\$defs/SubtypeForA"),
74+
)
75+
),
76+
"TypeForCompB" to JsonObject(mapOf("type" to JsonPrimitive("number"))),
77+
"SubtypeForA" to JsonObject(mapOf("type" to JsonPrimitive("boolean"))),
78+
)
79+
)
80+
)
81+
)
82+
83+
val catalogSchema =
84+
JsonObject(
85+
mapOf(
86+
A2uiConstants.CATALOG_ID_KEY to JsonPrimitive("basic"),
87+
A2uiConstants.CATALOG_COMPONENTS_KEY to
88+
JsonObject(
89+
mapOf(
90+
"CompA" to
91+
JsonObject(
92+
mapOf("\$ref" to JsonPrimitive("common_types.json#/\$defs/TypeForCompA"))
93+
),
94+
"CompB" to
95+
JsonObject(
96+
mapOf("\$ref" to JsonPrimitive("common_types.json#/\$defs/TypeForCompB"))
97+
),
98+
)
99+
),
100+
)
101+
)
102+
103+
val catalog =
104+
A2uiCatalog(
105+
version = A2uiVersion.VERSION_0_9,
106+
name = "test",
107+
serverToClientSchema = JsonObject(emptyMap()),
108+
commonTypesSchema = commonTypes,
109+
catalogSchema = catalogSchema,
110+
)
111+
112+
val prunedCatalog = catalog.withPrunedComponents(listOf("CompA"))
113+
val prunedDefs = prunedCatalog.commonTypesSchema["\$defs"]?.jsonObject
114+
115+
assertTrue(prunedDefs != null)
116+
assertTrue(prunedDefs.containsKey("TypeForCompA"))
117+
assertTrue(prunedDefs.containsKey("SubtypeForA"))
118+
assertFalse(prunedDefs.containsKey("TypeForCompB"))
119+
}
120+
}

0 commit comments

Comments
 (0)