Skip to content

Commit 6b9936b

Browse files
committed
WIP
1 parent e022209 commit 6b9936b

File tree

4 files changed

+177
-5
lines changed

4 files changed

+177
-5
lines changed

libs/core/kiln_ai/adapters/model_adapters/litellm_adapter.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -392,7 +392,7 @@ def json_schema_response_format(self) -> dict[str, Any]:
392392
raise ValueError(
393393
"Invalid output schema for this task. Cannot use JSON schema response format."
394394
)
395-
output_schema = close_object_schemas(output_schema)
395+
output_schema = close_object_schemas(output_schema, strict=True)
396396
return {
397397
"response_format": {
398398
"type": "json_schema",
@@ -410,7 +410,7 @@ def tool_call_params(self, strict: bool) -> dict[str, Any]:
410410
raise ValueError(
411411
"Invalid output schema for this task. Can not use tool calls."
412412
)
413-
output_schema = close_object_schemas(output_schema)
413+
output_schema = close_object_schemas(output_schema, strict=strict)
414414

415415
function_params = {
416416
"name": "task_response",

libs/core/kiln_ai/adapters/model_adapters/test_litellm_adapter.py

Lines changed: 77 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -195,7 +195,7 @@ async def test_response_format_options_json_schema(config, mock_task):
195195
patch.object(adapter, "has_structured_output", return_value=True),
196196
):
197197
options = await adapter.response_format_options()
198-
expected_schema = close_object_schemas(mock_task.output_schema())
198+
expected_schema = close_object_schemas(mock_task.output_schema(), strict=True)
199199
assert options == {
200200
"response_format": {
201201
"type": "json_schema",
@@ -235,7 +235,7 @@ def test_tool_call_params_strict(config, mock_task):
235235
adapter = LiteLlmAdapter(config=config, kiln_task=mock_task)
236236

237237
params = adapter.tool_call_params(strict=True)
238-
expected_schema = close_object_schemas(mock_task.output_schema())
238+
expected_schema = close_object_schemas(mock_task.output_schema(), strict=True)
239239

240240
assert params == {
241241
"tools": [
@@ -255,6 +255,81 @@ def test_tool_call_params_strict(config, mock_task):
255255
}
256256

257257

258+
def test_tool_call_params_strict_adds_required_to_nested(config, tmp_path):
259+
project_path = tmp_path / "test_project_nested" / "project.kiln"
260+
project_path.parent.mkdir()
261+
project = Project(name="Nested Project", path=str(project_path))
262+
project.save_to_file()
263+
264+
nested_schema = {
265+
"type": "object",
266+
"properties": {
267+
"user": {
268+
"type": "object",
269+
"properties": {
270+
"name": {"type": "string"},
271+
"age": {"type": "integer"},
272+
},
273+
},
274+
"status": {"type": "string"},
275+
},
276+
}
277+
task = Task(
278+
name="Nested Task",
279+
instruction="Test instruction",
280+
parent=project,
281+
output_json_schema=json.dumps(nested_schema),
282+
)
283+
task.save_to_file()
284+
285+
adapter = LiteLlmAdapter(config=config, kiln_task=task)
286+
params = adapter.tool_call_params(strict=True)
287+
288+
result_schema = params["tools"][0]["function"]["parameters"]
289+
assert result_schema["required"] == ["user", "status"]
290+
assert result_schema["properties"]["user"]["required"] == ["name", "age"]
291+
292+
293+
@pytest.mark.asyncio
294+
async def test_json_schema_response_format_adds_required_to_nested(config, tmp_path):
295+
project_path = tmp_path / "test_project_nested2" / "project.kiln"
296+
project_path.parent.mkdir()
297+
project = Project(name="Nested Project 2", path=str(project_path))
298+
project.save_to_file()
299+
300+
nested_schema = {
301+
"type": "object",
302+
"properties": {
303+
"result": {
304+
"type": "object",
305+
"properties": {
306+
"value": {"type": "number"},
307+
"unit": {"type": "string"},
308+
},
309+
},
310+
},
311+
}
312+
task = Task(
313+
name="Nested Task 2",
314+
instruction="Test instruction",
315+
parent=project,
316+
output_json_schema=json.dumps(nested_schema),
317+
)
318+
task.save_to_file()
319+
320+
config.run_config_properties.structured_output_mode = (
321+
StructuredOutputMode.json_schema
322+
)
323+
adapter = LiteLlmAdapter(config=config, kiln_task=task)
324+
325+
with patch.object(adapter, "has_structured_output", return_value=True):
326+
options = await adapter.response_format_options()
327+
328+
result_schema = options["response_format"]["json_schema"]["schema"]
329+
assert result_schema["required"] == ["result"]
330+
assert result_schema["properties"]["result"]["required"] == ["value", "unit"]
331+
332+
258333
@pytest.mark.parametrize(
259334
"provider_name,expected_prefix",
260335
[

libs/core/kiln_ai/datamodel/json_schema.py

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -157,14 +157,18 @@ def single_string_field_name(schema: Dict) -> str | None:
157157
return None
158158

159159

160-
def close_object_schemas(schema: Dict) -> Dict:
160+
def close_object_schemas(schema: Dict, strict: bool = False) -> Dict:
161161
"""Return a deep-copied schema with object nodes closed by default.
162162
163163
Any schema node with 'type == "object"' gets 'additionalProperties: false'
164164
if it is not already set. This normalization is recursive and walks nested
165165
schema structures such as 'properties', 'items', '$defs', and composed
166166
schemas like 'anyOf'/'oneOf'/'allOf'. Existing explicit
167167
'additionalProperties' values are preserved.
168+
169+
When strict=True, also sets 'required' to list all property keys on every
170+
object node with 'properties'. This is needed for OpenAI's strict
171+
structured output mode, which does not support optional properties.
168172
"""
169173

170174
def _normalize(node: Any) -> Any:
@@ -203,6 +207,9 @@ def _normalize(node: Any) -> Any:
203207
):
204208
normalized["additionalProperties"] = False
205209

210+
if strict and normalized.get("type") == "object" and "properties" in normalized:
211+
normalized["required"] = list(normalized["properties"].keys())
212+
206213
return normalized
207214

208215
return _normalize(schema)

libs/core/kiln_ai/datamodel/test_json_schema.py

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -248,6 +248,96 @@ def test_close_object_schemas_preserves_explicit_additional_properties():
248248
assert normalized["properties"]["metadata"]["additionalProperties"] is True
249249

250250

251+
def test_close_object_schemas_strict_adds_required():
252+
schema = {
253+
"type": "object",
254+
"properties": {
255+
"name": {"type": "string"},
256+
"age": {"type": "integer"},
257+
},
258+
}
259+
result = close_object_schemas(schema, strict=True)
260+
assert result["required"] == ["name", "age"]
261+
assert result["additionalProperties"] is False
262+
263+
264+
def test_close_object_schemas_strict_overwrites_partial_required():
265+
schema = {
266+
"type": "object",
267+
"properties": {
268+
"name": {"type": "string"},
269+
"age": {"type": "integer"},
270+
"email": {"type": "string"},
271+
},
272+
"required": ["name"],
273+
}
274+
result = close_object_schemas(schema, strict=True)
275+
assert result["required"] == ["name", "age", "email"]
276+
277+
278+
def test_close_object_schemas_strict_nested():
279+
schema = {
280+
"type": "object",
281+
"properties": {
282+
"user": {
283+
"type": "object",
284+
"properties": {
285+
"name": {"type": "string"},
286+
"address": {
287+
"type": "object",
288+
"properties": {
289+
"street": {"type": "string"},
290+
"city": {"type": "string"},
291+
},
292+
},
293+
},
294+
},
295+
"tags": {
296+
"type": "array",
297+
"items": {
298+
"type": "object",
299+
"properties": {
300+
"key": {"type": "string"},
301+
"value": {"type": "string"},
302+
},
303+
},
304+
},
305+
},
306+
}
307+
result = close_object_schemas(schema, strict=True)
308+
assert result["required"] == ["user", "tags"]
309+
assert result["properties"]["user"]["required"] == ["name", "address"]
310+
assert result["properties"]["user"]["properties"]["address"]["required"] == [
311+
"street",
312+
"city",
313+
]
314+
assert result["properties"]["tags"]["items"]["required"] == ["key", "value"]
315+
316+
317+
def test_close_object_schemas_non_strict_no_required():
318+
schema = {
319+
"type": "object",
320+
"properties": {
321+
"name": {"type": "string"},
322+
"age": {"type": "integer"},
323+
},
324+
}
325+
result = close_object_schemas(schema)
326+
assert "required" not in result
327+
328+
result_explicit = close_object_schemas(schema, strict=False)
329+
assert "required" not in result_explicit
330+
331+
332+
def test_close_object_schemas_strict_no_properties():
333+
schema = {
334+
"type": "object",
335+
"additionalProperties": {"type": "string"},
336+
}
337+
result = close_object_schemas(schema, strict=True)
338+
assert "required" not in result
339+
340+
251341
@pytest.mark.parametrize(
252342
"schema,expected",
253343
[

0 commit comments

Comments
 (0)