Skip to content

Commit 7f8b620

Browse files
PsiACEPsiACE
committed
refactor(codegen): name inline schema variants
1 parent 6f19a4a commit 7f8b620

5 files changed

Lines changed: 154 additions & 155 deletions

File tree

‎scripts/gen_schema.py‎

Lines changed: 87 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -34,22 +34,22 @@
3434
"integer+uint64=integer",
3535
)
3636

37-
# These are generator configuration, not schema rewrites. The v1.19 schema uses
38-
# discriminators for open variants that datamodel-code-generator otherwise renders
39-
# as placeholder literals.
40-
CATCHALL_TYPE_OVERRIDES = {
41-
"CreateElicitationRequest3.mode": "acp._schema_types.OpenString",
42-
"CreateElicitationRequest4.mode": "acp._schema_types.OpenString",
43-
"CreateElicitationResponse4.action": "acp._schema_types.OpenString",
44-
"ElicitationPropertySchema6.type": "acp._schema_types.OpenString",
45-
"MultiSelectItems2.type": "acp._schema_types.OpenString",
37+
OPEN_UNIONS = {
38+
"CreateElicitationRequest": 2,
39+
"CreateElicitationResponse": 3,
40+
"ElicitationPropertySchema": 5,
41+
"MultiSelectItems": 1,
4642
}
47-
OPEN_UNIONS = (
48-
"CreateElicitationRequest",
49-
"CreateElicitationResponse",
50-
"ElicitationPropertySchema",
51-
"MultiSelectItems",
52-
)
43+
44+
45+
def _inline_model_ref(definition: str, *steps: tuple[str, int | None]) -> str:
46+
ref = f"#/$defs/{definition}"
47+
for keyword, index in steps:
48+
ref += f"#-datamodel-code-generator-#-{keyword}-#-special-#"
49+
if index is not None:
50+
ref += f"/{index}"
51+
return ref
52+
5353

5454
# A few schema definitions share their name with a tagged SessionUpdate or
5555
# MultiSelectItems variant. Give the definition an internal base name so the public
@@ -65,6 +65,72 @@
6565
"UsageUpdate",
6666
)
6767
}
68+
MODEL_NAME_MAP.update({
69+
_inline_model_ref("AgentResponse", ("anyOf", 0), ("object", None)): "AgentResponseMessage",
70+
_inline_model_ref("AgentResponse", ("anyOf", 1), ("object", None)): "AgentErrorMessage",
71+
_inline_model_ref("ClientResponse", ("anyOf", 0), ("object", None)): "ClientResponseMessage",
72+
_inline_model_ref("ClientResponse", ("anyOf", 1), ("object", None)): "ClientErrorMessage",
73+
_inline_model_ref("SetSessionConfigOptionRequest", ("anyOf", 0), ("object", None)): (
74+
"SetSessionConfigOptionBooleanRequest"
75+
),
76+
_inline_model_ref("SetSessionConfigOptionRequest", ("anyOf", 1), ("object", None)): (
77+
"SetSessionConfigOptionSelectRequest"
78+
),
79+
_inline_model_ref("CreateElicitationResponse", ("anyOf", 0), ("allOf", None)): "AcceptElicitationResponse",
80+
_inline_model_ref("CreateElicitationResponse", ("anyOf", 1), ("object", None)): "DeclineElicitationResponse",
81+
_inline_model_ref("CreateElicitationResponse", ("anyOf", 2), ("object", None)): "CancelElicitationResponse",
82+
_inline_model_ref("CreateElicitationResponse", ("anyOf", 3), ("object", None)): "OtherElicitationResponse",
83+
_inline_model_ref("ElicitationFormMode", ("anyOf", 0), ("allOf", None)): "ElicitationFormSessionMode",
84+
_inline_model_ref("ElicitationFormMode", ("anyOf", 1), ("allOf", None)): "ElicitationFormRequestMode",
85+
_inline_model_ref("ElicitationUrlMode", ("anyOf", 0), ("allOf", None)): "ElicitationUrlSessionMode",
86+
_inline_model_ref("ElicitationUrlMode", ("anyOf", 1), ("allOf", None)): "ElicitationUrlRequestMode",
87+
_inline_model_ref("ElicitationPropertySchema", ("anyOf", 0), ("allOf", None)): ("ElicitationStringPropertySchema"),
88+
_inline_model_ref("ElicitationPropertySchema", ("anyOf", 1), ("allOf", None)): ("ElicitationNumberPropertySchema"),
89+
_inline_model_ref("ElicitationPropertySchema", ("anyOf", 2), ("allOf", None)): ("ElicitationIntegerPropertySchema"),
90+
_inline_model_ref("ElicitationPropertySchema", ("anyOf", 3), ("allOf", None)): ("ElicitationBooleanPropertySchema"),
91+
_inline_model_ref("ElicitationPropertySchema", ("anyOf", 4), ("allOf", None)): (
92+
"ElicitationMultiSelectPropertySchema"
93+
),
94+
_inline_model_ref("ElicitationPropertySchema", ("anyOf", 5), ("object", None)): ("ElicitationOtherPropertySchema"),
95+
_inline_model_ref("MultiSelectItems", ("anyOf", 0), ("allOf", None)): "StringMultiSelectItems",
96+
_inline_model_ref("MultiSelectItems", ("anyOf", 1), ("object", None)): "OtherMultiSelectItems",
97+
_inline_model_ref("CreateElicitationRequest", ("anyOf", 0), ("allOf", None), ("allOf", None)): (
98+
"CreateFormElicitationRequestBase"
99+
),
100+
_inline_model_ref("CreateElicitationRequest", ("anyOf", 0), ("allOf", 0), ("allOf", None)): (
101+
"CreateFormSessionElicitationRequestBase"
102+
),
103+
_inline_model_ref("CreateElicitationRequest", ("anyOf", 0), ("allOf", 1), ("allOf", None)): (
104+
"CreateFormRequestElicitationRequestBase"
105+
),
106+
_inline_model_ref("CreateElicitationRequest", ("anyOf", 0), ("allOf", None), ("union_model-0", None)): (
107+
"CreateFormSessionElicitationRequest"
108+
),
109+
_inline_model_ref("CreateElicitationRequest", ("anyOf", 0), ("allOf", None), ("union_model-1", None)): (
110+
"CreateFormRequestElicitationRequest"
111+
),
112+
_inline_model_ref("CreateElicitationRequest", ("anyOf", 1), ("allOf", None), ("allOf", None)): (
113+
"CreateUrlElicitationRequestBase"
114+
),
115+
_inline_model_ref("CreateElicitationRequest", ("anyOf", 1), ("allOf", 0), ("allOf", None)): (
116+
"CreateUrlSessionElicitationRequestBase"
117+
),
118+
_inline_model_ref("CreateElicitationRequest", ("anyOf", 1), ("allOf", 1), ("allOf", None)): (
119+
"CreateUrlRequestElicitationRequestBase"
120+
),
121+
_inline_model_ref("CreateElicitationRequest", ("anyOf", 1), ("allOf", None), ("union_model-0", None)): (
122+
"CreateUrlSessionElicitationRequest"
123+
),
124+
_inline_model_ref("CreateElicitationRequest", ("anyOf", 1), ("allOf", None), ("union_model-1", None)): (
125+
"CreateUrlRequestElicitationRequest"
126+
),
127+
_inline_model_ref("CreateElicitationRequest", ("anyOf", 2), ("anyOf", 0), ("allOf", None)): (
128+
"CreateOtherSessionElicitationRequest"
129+
),
130+
_inline_model_ref("CreateElicitationRequest", ("anyOf", 2), ("anyOf", 1), ("allOf", None)): (
131+
"CreateOtherRequestElicitationRequest"
132+
),
133+
})
68134

69135
# datamodel-code-generator owns schema interpretation and its internal model names.
70136
# This block only preserves the Python names already published by the SDK.
@@ -87,11 +153,6 @@
87153
"other",
88154
]
89155
90-
AgentResponseMessage = AgentResponse1
91-
AgentErrorMessage = AgentResponse2
92-
ClientResponseMessage = ClientResponse1
93-
ClientErrorMessage = ClientResponse2
94-
95156
TextContentBlock = ContentBlockText
96157
ImageContentBlock = ContentBlockImage
97158
AudioContentBlock = ContentBlockAudio
@@ -105,11 +166,6 @@
105166
AllowedOutcome = RequestPermissionOutcomeSelected
106167
EnvVarAuthMethod = AuthMethodEnvVarModel
107168
TerminalAuthMethod = AuthMethodTerminalModel
108-
SetSessionConfigOptionBooleanRequest = SetSessionConfigOptionRequest1
109-
SetSessionConfigOptionSelectRequest = SetSessionConfigOptionRequest2
110-
SetSessionConfigOptionRequest1.__acp_public_name__ = "SetSessionConfigOptionBooleanRequest"
111-
SetSessionConfigOptionRequest2.__acp_public_name__ = "SetSessionConfigOptionSelectRequest"
112-
113169
UserMessageChunk = SessionUpdateUserMessageChunk
114170
AgentMessageChunk = SessionUpdateAgentMessageChunk
115171
AgentThoughtChunk = SessionUpdateAgentThoughtChunk
@@ -138,13 +194,9 @@
138194
TerminalToolCallContent = ToolCallContentTerminal
139195
140196
CreateOtherElicitationRequest = Union[
141-
CreateElicitationRequest3,
142-
CreateElicitationRequest4,
197+
CreateOtherSessionElicitationRequest,
198+
CreateOtherRequestElicitationRequest,
143199
]
144-
CreateFormSessionElicitationRequest = CreateElicitationRequest14
145-
CreateFormRequestElicitationRequest = CreateElicitationRequest15
146-
CreateUrlSessionElicitationRequest = CreateElicitationRequest24
147-
CreateUrlRequestElicitationRequest = CreateElicitationRequest25
148200
CreateFormElicitationRequest = Union[
149201
CreateFormSessionElicitationRequest,
150202
CreateFormRequestElicitationRequest,
@@ -159,36 +211,20 @@
159211
CreateOtherElicitationRequest,
160212
]
161213
162-
AcceptElicitationResponse = CreateElicitationResponse1
163-
DeclineElicitationResponse = CreateElicitationResponse2
164-
CancelElicitationResponse = CreateElicitationResponse3
165-
OtherElicitationResponse = CreateElicitationResponse4
166214
CreateElicitationResponse = Union[
167215
AcceptElicitationResponse,
168216
DeclineElicitationResponse,
169217
CancelElicitationResponse,
170218
OtherElicitationResponse,
171219
]
172-
ElicitationFormSessionMode = ElicitationFormMode1
173-
ElicitationFormRequestMode = ElicitationFormMode2
174-
ElicitationUrlSessionMode = ElicitationUrlMode1
175-
ElicitationUrlRequestMode = ElicitationUrlMode2
176220
ElicitationMode = Union[
177221
ElicitationFormSessionMode,
178222
ElicitationFormRequestMode,
179223
ElicitationUrlSessionMode,
180224
ElicitationUrlRequestMode,
181225
]
182226
183-
ElicitationStringPropertySchema = ElicitationPropertySchema1
184-
ElicitationNumberPropertySchema = ElicitationPropertySchema2
185-
ElicitationIntegerPropertySchema = ElicitationPropertySchema3
186-
ElicitationBooleanPropertySchema = ElicitationPropertySchema4
187-
ElicitationMultiSelectPropertySchema = ElicitationPropertySchema5
188-
ElicitationOtherPropertySchema = ElicitationPropertySchema6
189227
_StringMultiSelectItems = StringMultiSelectItemsBase
190-
StringMultiSelectItems = MultiSelectItems1
191-
OtherMultiSelectItems = MultiSelectItems2
192228
193229
NesEditSuggestionVariant = NesSuggestionEdit
194230
NesJumpSuggestionVariant = NesSuggestionJump
@@ -264,7 +300,6 @@ def render_schema() -> str:
264300
schema_version="2020-12",
265301
schema_version_mode=VersionMode.Strict,
266302
type_mappings=list(UNSIGNED_TYPE_MAPPINGS),
267-
type_overrides=CATCHALL_TYPE_OVERRIDES,
268303
generate_schema_validators=True,
269304
use_annotated=True,
270305
field_constraints=True,
@@ -277,13 +312,14 @@ def render_schema() -> str:
277312

278313

279314
def _schema_for_codegen(schema: dict[str, Any]) -> dict[str, Any]:
280-
"""Apply the fixed v1.19 compatibility patch required for open unions."""
315+
"""Drop open-union constraints that Pydantic cannot represent statically."""
281316
patched = copy.deepcopy(schema)
282-
for name in OPEN_UNIONS:
317+
for name, catchall_index in OPEN_UNIONS.items():
283318
try:
284319
del patched["$defs"][name]["discriminator"]
320+
del patched["$defs"][name]["anyOf"][catchall_index]["not"]
285321
except KeyError:
286-
raise ValueError(f"{name} no longer has the expected discriminator") from None
322+
raise ValueError(f"{name} no longer has the expected open-union shape") from None
287323
return patched
288324

289325

‎src/acp/_schema_base.py‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -20,14 +20,14 @@ class BaseModel(pydantic.BaseModel):
2020
"""Runtime behavior shared by generated ACP schema models."""
2121

2222
_reserved_tags: ClassVar[dict[str, tuple[str, frozenset[str]]]] = {
23-
"CreateElicitationRequest3": ("mode", frozenset({"form", "url"})),
24-
"CreateElicitationRequest4": ("mode", frozenset({"form", "url"})),
25-
"CreateElicitationResponse4": ("action", frozenset({"accept", "cancel", "decline"})),
26-
"ElicitationPropertySchema6": (
23+
"CreateOtherSessionElicitationRequest": ("mode", frozenset({"form", "url"})),
24+
"CreateOtherRequestElicitationRequest": ("mode", frozenset({"form", "url"})),
25+
"OtherElicitationResponse": ("action", frozenset({"accept", "cancel", "decline"})),
26+
"ElicitationOtherPropertySchema": (
2727
"type",
2828
frozenset({"array", "boolean", "integer", "number", "string"}),
2929
),
30-
"MultiSelectItems2": ("type", frozenset({"string"})),
30+
"OtherMultiSelectItems": ("type", frozenset({"string"})),
3131
}
3232

3333
model_config = ConfigDict(

‎src/acp/_schema_types.py‎

Lines changed: 0 additions & 3 deletions
This file was deleted.

0 commit comments

Comments
 (0)