From c4105802515dd3f28ce24fc770475a994d4c5aec Mon Sep 17 00:00:00 2001 From: -LAN- Date: Mon, 3 Aug 2026 10:05:09 +0800 Subject: [PATCH 1/3] fix(llm): enforce reliable structured output Normalize the legacy structured-output switch to the canonical field, validate Draft 7 schemas and final results, and preserve provider capability tri-state behavior. --- src/graphon/dsl/slim/package_loader.py | 5 +- .../model_runtime/entities/model_entities.py | 9 +- src/graphon/nodes/llm/entities.py | 15 +- src/graphon/nodes/llm/node.py | 97 ++++++--- tests/dsl/test_slim_llm.py | 42 +++- tests/dsl/test_slim_package_loader.py | 45 +++++ tests/nodes/llm/test_node.py | 188 ++++++++++++++++++ 7 files changed, 360 insertions(+), 41 deletions(-) diff --git a/src/graphon/dsl/slim/package_loader.py b/src/graphon/dsl/slim/package_loader.py index fc9359c3..a7aee1fd 100644 --- a/src/graphon/dsl/slim/package_loader.py +++ b/src/graphon/dsl/slim/package_loader.py @@ -284,9 +284,10 @@ def _convert_model_entity(self, raw_model: dict[str, Any]) -> AIModelEntity | No if model_type is None or fetch_from is None: return None + raw_features = raw_model.get("features") features = [ feature - for item in raw_model.get("features", []) or [] + for item in raw_features or [] if (feature := self._convert_model_feature(item)) is not None ] model_properties = { @@ -299,7 +300,7 @@ def _convert_model_entity(self, raw_model: dict[str, Any]) -> AIModelEntity | No model=str(raw_model["model"]), label=self._convert_i18n(raw_model.get("label")), model_type=model_type, - features=features or None, + features=features if raw_features is not None else None, fetch_from=fetch_from, model_properties=model_properties, deprecated=bool(raw_model.get("deprecated")), diff --git a/src/graphon/model_runtime/entities/model_entities.py b/src/graphon/model_runtime/entities/model_entities.py index ccda57c2..add0c93c 100644 --- a/src/graphon/model_runtime/entities/model_entities.py +++ b/src/graphon/model_runtime/entities/model_entities.py @@ -200,12 +200,11 @@ def validate_model(self) -> Self: ), None, ) - if not schema_key: + # Explicit feature lists are authoritative; infer support only for legacy + # model schemas that omit the feature declaration. + if not schema_key or self.features is not None: return self - if self.features is None: - self.features = [ModelFeature.STRUCTURED_OUTPUT] - elif ModelFeature.STRUCTURED_OUTPUT not in self.features: - self.features.append(ModelFeature.STRUCTURED_OUTPUT) + self.features = [ModelFeature.STRUCTURED_OUTPUT] return self diff --git a/src/graphon/nodes/llm/entities.py b/src/graphon/nodes/llm/entities.py index 18e233bf..133e18de 100644 --- a/src/graphon/nodes/llm/entities.py +++ b/src/graphon/nodes/llm/entities.py @@ -1,7 +1,7 @@ from collections.abc import Mapping, Sequence from typing import Any, Literal -from pydantic import BaseModel, Field, field_validator +from pydantic import BaseModel, Field, field_validator, model_validator from graphon.entities.base_node_data import BaseNodeData from graphon.enums import BuiltinNodeTypes, NodeType @@ -75,8 +75,7 @@ class LLMNodeData(BaseNodeData): context: ContextConfig vision: VisionConfig = Field(default_factory=VisionConfig) structured_output: Mapping[str, Any] | None = None - # We used 'structured_output_enabled' in the past, but it's not a good name. - structured_output_switch_on: bool = Field(False, alias="structured_output_enabled") + structured_output_switch_on: bool = False reasoning_format: Literal["separated", "tagged"] = Field( # Keep tagged as default for backward compatibility default="tagged", @@ -96,6 +95,16 @@ class LLMNodeData(BaseNodeData): ), ) + @model_validator(mode="before") + @classmethod + def migrate_legacy_structured_output_switch(cls, data: Any) -> Any: + if not isinstance(data, Mapping) or "structured_output_enabled" not in data: + return data + data = dict(data) + legacy_value = data.pop("structured_output_enabled") + data.setdefault("structured_output_switch_on", legacy_value) + return data + @field_validator("prompt_config", mode="before") @classmethod def convert_none_prompt_config(cls, v: Any) -> Any: diff --git a/src/graphon/nodes/llm/node.py b/src/graphon/nodes/llm/node.py index 137c010a..e99c2626 100644 --- a/src/graphon/nodes/llm/node.py +++ b/src/graphon/nodes/llm/node.py @@ -11,6 +11,8 @@ from datetime import UTC, datetime, timedelta from typing import Any, Literal, assert_never, override +from jsonschema import Draft7Validator, SchemaError, ValidationError + from graphon.entities.graph_init_params import GraphInitParams from graphon.enums import ( BuiltinNodeTypes, @@ -44,7 +46,7 @@ TextPromptMessageContent, UserPromptMessage, ) -from graphon.model_runtime.entities.model_entities import ModelPropertyKey +from graphon.model_runtime.entities.model_entities import ModelFeature, ModelPropertyKey from graphon.model_runtime.memory.prompt_message_memory import PromptMessageMemory from graphon.model_runtime.utils.encoders import jsonable_encoder from graphon.node_events.base import ( @@ -296,6 +298,7 @@ def _prepare_run_prompt( context=collected_context.context or "", memory=self._memory, model_instance=model_instance, + structured_output_enabled=self.node_data.structured_output_enabled, stop=model_instance.stop, prompt_template=self.node_data.prompt_template, memory_config=self.node_data.memory, @@ -404,7 +407,7 @@ def _yield_run_completion( finish_reason = event.finish_reason reasoning_content = event.reasoning_content or "" clean_text = self._extract_clean_text(event.text) - if event.structured_output: + if event.structured_output is not None: structured_output = LLMStructuredOutput( structured_output=event.structured_output, ) @@ -538,6 +541,7 @@ def _invoke_llm_with_polling( model_instance=self._model_instance, reasoning_format=self.node_data.reasoning_format, request_start_time=request_start_time, + json_schema=json_schema, ) return case LLMPollingStatus.FAILED: @@ -755,6 +759,7 @@ def invoke_llm( model_parameters = model_instance.parameters invoke_model_parameters = dict(model_parameters) invoke_result: LLMResult | Generator[LLMResultChunk, None, None] + output_schema: dict[str, Any] | None = None if structured_output_enabled: output_schema = LLMNode.fetch_structured_output_schema( structured_output=structured_output or {}, @@ -787,6 +792,7 @@ def invoke_llm( model_instance=model_instance, reasoning_format=reasoning_format, request_start_time=request_start_time, + json_schema=output_schema, ) @staticmethod @@ -800,6 +806,7 @@ def handle_invoke_result( model_instance: LLMProtocol, reasoning_format: Literal["separated", "tagged"] = "tagged", request_start_time: float | None = None, + json_schema: Mapping[str, Any] | None = None, ) -> Generator[NodeEventBase | LLMStructuredOutput, None, None]: if isinstance(invoke_result, LLMResult): yield from LLMNode._yield_blocking_invoke_result( @@ -808,6 +815,7 @@ def handle_invoke_result( file_outputs=file_outputs, reasoning_format=reasoning_format, request_start_time=request_start_time, + json_schema=json_schema, ) return @@ -819,6 +827,7 @@ def handle_invoke_result( model_instance=model_instance, reasoning_format=reasoning_format, request_start_time=request_start_time, + json_schema=json_schema, ) @staticmethod @@ -829,19 +838,25 @@ def _yield_blocking_invoke_result( file_outputs: list[File], reasoning_format: Literal["separated", "tagged"] = "tagged", request_start_time: float | None = None, + json_schema: Mapping[str, Any] | None = None, ) -> Generator[ModelInvokeCompletedEvent, None, None]: duration = None if request_start_time is not None: duration = time.perf_counter() - request_start_time invoke_result.usage.latency = round(duration, 3) - yield LLMNode.handle_blocking_result( + event = LLMNode.handle_blocking_result( invoke_result=invoke_result, saver=file_saver, file_outputs=file_outputs, reasoning_format=reasoning_format, request_latency=duration, ) + LLMNode._validate_structured_output_result( + structured_output=event.structured_output, + json_schema=json_schema, + ) + yield event @staticmethod def _yield_streaming_invoke_result( @@ -853,6 +868,7 @@ def _yield_streaming_invoke_result( model_instance: LLMProtocol, reasoning_format: Literal["separated", "tagged"] = "tagged", request_start_time: float | None = None, + json_schema: Mapping[str, Any] | None = None, ) -> Generator[NodeEventBase | LLMStructuredOutput, None, None]: start_time = ( request_start_time @@ -876,10 +892,10 @@ def _yield_streaming_invoke_result( model_instance=model_instance, error=e, ): - msg = f"Failed to parse structured output: {e}" + msg = f"Failed to parse structured output (stage=result, path=$): {e}" raise LLMNodeError(msg) from e if type(e).__name__ == "OutputParserError": - msg = f"Failed to parse structured output: {e}" + msg = f"Failed to parse structured output (stage=result, path=$): {e}" raise LLMNodeError(msg) from e raise @@ -917,6 +933,10 @@ def _yield_streaming_invoke_result( first_token_time=state.first_token_time, start_time=state.start_time, ) + LLMNode._validate_structured_output_result( + structured_output=state.structured_output, + json_schema=json_schema, + ) yield ModelInvokeCompletedEvent( # Use clean_text for separated mode, full_text for tagged mode @@ -1080,6 +1100,29 @@ def _is_structured_output_parse_error( and is_structured_output_parse_error(error) ) + @staticmethod + def _validate_structured_output_result( + *, + structured_output: Mapping[str, Any] | None, + json_schema: Mapping[str, Any] | None, + ) -> None: + if json_schema is None: + return + if structured_output is None: + msg = ( + "Structured output validation failed " + "(stage=result, path=$): structured output is missing" + ) + raise LLMNodeError(msg) + try: + Draft7Validator(json_schema).validate(structured_output) + except ValidationError as error: + msg = ( + "Structured output validation failed " + f"(stage=result, path={error.json_path}): {error.message}" + ) + raise LLMNodeError(msg) from error + @staticmethod def _finalize_streaming_usage( *, @@ -1353,6 +1396,7 @@ def fetch_prompt_messages( context: str = "", memory: PromptMessageMemory | None = None, model_instance: LLMProtocol, + structured_output_enabled: bool = False, prompt_template: Sequence[LLMNodeChatModelMessage] | LLMNodeCompletionModelPromptTemplate, stop: Sequence[str] | None = None, @@ -1365,6 +1409,17 @@ def fetch_prompt_messages( jinja2_template_renderer: Jinja2TemplateRenderer | None = None, ) -> tuple[Sequence[PromptMessage], Sequence[str] | None]: model_schema = llm_utils.fetch_model_schema(model_instance=model_instance) + if ( + structured_output_enabled + and model_schema.features is not None + and ModelFeature.STRUCTURED_OUTPUT not in model_schema.features + ): + msg = ( + "Structured output is not supported by model " + f"{model_instance.provider}/{model_instance.model_name} " + "(stage=capability)" + ) + raise LLMNodeError(msg) prompt_messages = LLMNode._build_prompt_messages_from_template( sys_query=sys_query, context=context, @@ -2073,27 +2128,23 @@ def fetch_structured_output_schema( or not a JSON object. """ - if not structured_output: - msg = "Please provide a valid structured output schema" - raise LLMNodeError(msg) - structured_output_schema = json.dumps( - structured_output.get("schema", {}), - ensure_ascii=False, - ) - if not structured_output_schema: - msg = "Please provide a valid structured output schema" + raw_schema = structured_output.get("schema") + if not isinstance(raw_schema, Mapping): + msg = ( + "Invalid structured output schema " + "(stage=schema, path=$.schema): expected a JSON object" + ) raise LLMNodeError(msg) - + schema = dict(raw_schema) try: - schema = json.loads(structured_output_schema) - if not isinstance(schema, dict): - msg = "structured_output_schema must be a JSON object" - raise LLMNodeError(msg) - except json.JSONDecodeError as error: - msg = "structured_output_schema is not valid JSON format" + Draft7Validator.check_schema(schema) + except SchemaError as error: + msg = ( + "Invalid structured output schema " + f"(stage=schema, path={error.json_path}): {error.message}" + ) raise LLMNodeError(msg) from error - else: - return schema + return schema @staticmethod def _save_multimodal_output_and_convert_result_to_markdown( diff --git a/tests/dsl/test_slim_llm.py b/tests/dsl/test_slim_llm.py index e6b4decb..3d7ae1b0 100644 --- a/tests/dsl/test_slim_llm.py +++ b/tests/dsl/test_slim_llm.py @@ -1,5 +1,6 @@ from __future__ import annotations +import json from collections.abc import Iterable, Mapping from pathlib import Path from typing import Any @@ -26,14 +27,16 @@ def invoke_chunks( self.calls.append((plugin_id, action, data)) if action == "get_llm_num_tokens": return [{"num_tokens": 7}] - return [ - { - "delta": { - "index": 0, - "message": {"content": "hello"}, - } - } - ] + chunk = { + "delta": { + "index": 0, + "message": {"content": "hello"}, + }, + } + model_parameters = data.get("model_parameters") + if isinstance(model_parameters, Mapping) and "json_schema" in model_parameters: + chunk["structured_output"] = {"ok": True} + return [chunk] class _FailingSlimClient: @@ -179,6 +182,29 @@ def test_slim_llm_counts_tokens_and_collects_blocking_result( } +def test_slim_llm_passes_merged_parameters_and_json_schema( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + client = _patch_recording_slim_client(monkeypatch) + llm = _build_llm(tmp_path) + schema = {"type": "object", "required": ["ok"]} + + result = llm.invoke_llm_with_structured_output( + prompt_messages=[], + json_schema=schema, + model_parameters={"max_tokens": 8}, + stop=None, + stream=False, + ) + + assert result.structured_output == {"ok": True} + model_parameters = client.calls[-1][2]["model_parameters"] + assert model_parameters["temperature"] == pytest.approx(0.2) + assert model_parameters["max_tokens"] == 8 + assert json.loads(model_parameters["json_schema"]) == schema + + def test_slim_llm_preserves_slim_client_errors( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, diff --git a/tests/dsl/test_slim_package_loader.py b/tests/dsl/test_slim_package_loader.py index 89d0b6c1..260c036f 100644 --- a/tests/dsl/test_slim_package_loader.py +++ b/tests/dsl/test_slim_package_loader.py @@ -60,6 +60,51 @@ def test_slim_config_auto_discovers_uv_and_python( assert config.local.uv_path == "/usr/local/bin/uv" +@pytest.mark.parametrize( + ("raw_features", "infer_from_rule", "expected_features"), + [ + (None, False, None), + ([], True, []), + (None, True, [ModelFeature.STRUCTURED_OUTPUT]), + ( + ["structured-output"], + False, + [ModelFeature.STRUCTURED_OUTPUT], + ), + ], + ids=["unknown", "unsupported", "inferred", "supported"], +) +def test_slim_package_loader_preserves_model_feature_tri_state( + tmp_path: Path, + raw_features: list[str] | None, + infer_from_rule: bool, + expected_features: list[ModelFeature] | None, +) -> None: + loader = SlimPackageLoader( + SlimConfig( + bindings=[SlimProviderBinding(plugin_id="author/fake:0.0.1@test")], + local=SlimLocalSettings(folder=tmp_path), + ), + ) + raw_model = { + "model": "chat-model", + "label": {"en_US": "Chat Model"}, + "model_type": "llm", + "fetch_from": "predefined-model", + "model_properties": {}, + "parameter_rules": ( + [{"name": "json_schema", "type": "string"}] if infer_from_rule else [] + ), + } + if raw_features is not None: + raw_model["features"] = raw_features + + model = loader.convert_model_entity(raw_model) + + assert model is not None + assert model.features == expected_features + + def _write_multi_provider_plugin(plugin_root: Path) -> None: (plugin_root / "_assets").mkdir(parents=True, exist_ok=True) (plugin_root / "provider").mkdir(parents=True, exist_ok=True) diff --git a/tests/nodes/llm/test_node.py b/tests/nodes/llm/test_node.py index 043b237f..2d693565 100644 --- a/tests/nodes/llm/test_node.py +++ b/tests/nodes/llm/test_node.py @@ -8,6 +8,7 @@ import pytest +from graphon.entities.base_node_data import BaseNodeData from graphon.enums import WorkflowNodeExecutionStatus from graphon.file import helpers as file_helpers from graphon.file.enums import FileTransferMethod, FileType @@ -24,6 +25,7 @@ LLMResultChunk, LLMResultChunkDelta, LLMResultChunkWithStructuredOutput, + LLMResultWithStructuredOutput, LLMStructuredOutput, LLMUsage, ) @@ -188,6 +190,192 @@ def _stub_simple_prompt(monkeypatch: pytest.MonkeyPatch, node: LLMNode) -> None: ) +@pytest.mark.parametrize( + "switch_values", + [ + {"structured_output_enabled": True}, + {"structured_output_switch_on": True}, + { + "structured_output_enabled": False, + "structured_output_switch_on": True, + }, + ], + ids=["legacy", "current", "current-wins-conflict"], +) +def test_structured_output_switch_survives_node_data_round_trip( + switch_values: dict[str, bool], +) -> None: + payload = { + "type": "llm", + "model": { + "provider": "openai", + "name": "gpt-4o", + "mode": "chat", + }, + "prompt_template": [{"role": "user", "text": "Hello"}], + "context": {"enabled": False}, + "structured_output": {"schema": {"type": "object"}}, + **switch_values, + } + + node_data = LLMNode.validate_node_data(BaseNodeData.model_validate(payload)) + restored = LLMNode.validate_node_data(node_data) + dumped = restored.model_dump(mode="python", by_alias=True) + restored_from_dump = LLMNode.validate_node_data(dumped) + + assert restored.structured_output_switch_on is True + assert restored.structured_output_enabled is True + assert dumped["structured_output_switch_on"] is True + assert "structured_output_enabled" not in dumped + assert restored_from_dump.structured_output_switch_on is True + + +def test_fetch_structured_output_schema_checks_draft7_schema() -> None: + with pytest.raises(LLMNodeError) as exc_info: + LLMNode.fetch_structured_output_schema( + structured_output={"schema": {"type": "not-a-json-schema-type"}}, + ) + + assert "stage=schema" in str(exc_info.value) + assert "path=$.type" in str(exc_info.value) + + +_RESULT_SCHEMA = { + "type": "object", + "properties": { + "profile": { + "type": "object", + "properties": { + "role": {"type": "string", "enum": ["admin"]}, + "scores": {"type": "array", "items": {"type": "integer"}}, + }, + "required": ["role", "scores"], + }, + }, + "required": ["profile"], +} + + +@pytest.mark.parametrize( + ("structured_output", "path"), + [ + (None, "$"), + ({}, "$"), + ({"profile": {}}, "$.profile"), + ({"profile": {"role": "viewer", "scores": [1]}}, "$.profile.role"), + ( + {"profile": {"role": "admin", "scores": [1, "invalid"]}}, + "$.profile.scores[1]", + ), + ], +) +def test_final_structured_output_is_validated_against_full_schema( + structured_output: dict[str, Any] | None, + path: str, +) -> None: + chunk = LLMResultChunkWithStructuredOutput( + model="gpt-4o", + delta=LLMResultChunkDelta( + index=0, + message=AssistantPromptMessage(content=""), + usage=LLMUsage.empty_usage(), + ), + structured_output=structured_output, + ) + model = MagicMock(is_structured_output_parse_error=lambda _error: False) + + with pytest.raises(LLMNodeError) as exc_info: + list( + LLMNode.handle_invoke_result( + invoke_result=_stream_results(chunk), + file_saver=MagicMock(), + file_outputs=[], + node_id="llm", + model_instance=cast(LLMProtocol, model), + json_schema=_RESULT_SCHEMA, + ), + ) + + assert "stage=result" in str(exc_info.value) + assert f"path={path}" in str(exc_info.value) + if structured_output is None: + assert "structured output is missing" in str(exc_info.value) + + +def test_blocking_structured_output_is_validated_against_schema() -> None: + result = LLMResultWithStructuredOutput( + model="gpt-4o", + message=AssistantPromptMessage(content=""), + usage=LLMUsage.empty_usage(), + structured_output={"profile": {"role": "admin", "scores": ["invalid"]}}, + ) + + with pytest.raises(LLMNodeError, match=r"stage=result, path=\$\.profile\.scores"): + list( + LLMNode.handle_invoke_result( + invoke_result=result, + file_saver=MagicMock(), + file_outputs=[], + node_id="llm", + model_instance=cast(LLMProtocol, MagicMock()), + json_schema=_RESULT_SCHEMA, + ), + ) + + +@pytest.mark.parametrize( + ("features", "should_invoke"), + [ + ([ModelFeature.STRUCTURED_OUTPUT], True), + ([], False), + (None, True), + ], + ids=["supported", "unsupported", "unknown"], +) +def test_structured_output_capability_is_tri_state( + features: list[ModelFeature] | None, + should_invoke: bool, +) -> None: + model = MagicMock( + provider="openai", + model_name="gpt-4o", + parameters={}, + stop=(), + is_structured_output_parse_error=lambda _error: False, + ) + model.get_model_schema.return_value = SimpleNamespace( + features=features, + supports_prompt_content_type=lambda _content_type: True, + ) + model.invoke_llm_with_structured_output.return_value = _stream_results( + LLMResultChunkWithStructuredOutput( + model="gpt-4o", + delta=LLMResultChunkDelta( + index=0, + message=AssistantPromptMessage(content=""), + usage=LLMUsage.empty_usage(), + ), + structured_output={}, + ), + ) + node = _build_llm_node(model_instance=model) + node.node_data.structured_output_switch_on = True + node.node_data.structured_output = {"schema": {"type": "object"}} + + completed = next( + event for event in node._run() if isinstance(event, StreamCompletedEvent) + ) + + if should_invoke: + model.invoke_llm_with_structured_output.assert_called_once() + assert completed.node_run_result.status == WorkflowNodeExecutionStatus.SUCCEEDED + assert completed.node_run_result.outputs["structured_output"] == {} + else: + model.invoke_llm_with_structured_output.assert_not_called() + assert completed.node_run_result.status == WorkflowNodeExecutionStatus.FAILED + assert "stage=capability" in completed.node_run_result.error + + def test_run_emits_model_identity_in_node_result_inputs( monkeypatch: pytest.MonkeyPatch, ) -> None: From 3fc50c6c7ee0f9e3e32fdca38c88a46b76c59c34 Mon Sep 17 00:00:00 2001 From: -LAN- Date: Fri, 7 Aug 2026 09:15:21 +0800 Subject: [PATCH 2/3] fix(llm): reject structured output schema refs Fail schema validation before provider invocation so user-controlled references cannot trigger retrieval. --- src/graphon/nodes/llm/node.py | 37 ++++++++++++++++++++++++++++ tests/nodes/llm/test_node.py | 46 +++++++++++++++++++++++++++++++++++ 2 files changed, 83 insertions(+) diff --git a/src/graphon/nodes/llm/node.py b/src/graphon/nodes/llm/node.py index d5763d8b..3bec6c19 100644 --- a/src/graphon/nodes/llm/node.py +++ b/src/graphon/nodes/llm/node.py @@ -104,6 +104,37 @@ PromptMessageContentType.AUDIO: FileType.AUDIO, PromptMessageContentType.DOCUMENT: FileType.DOCUMENT, } +_DRAFT7_LITERAL_KEYWORDS = frozenset({"const", "default", "enum", "examples"}) + + +def _json_path_child(path: str, key: str | int) -> str: + if isinstance(key, int): + return f"{path}[{key}]" + if key.isidentifier(): + return f"{path}.{key}" + return f"{path}[{key!r}]" + + +def _find_draft7_ref_path(schema: Mapping[str, Any]) -> str | None: + pending: list[tuple[str, Any]] = [("$", schema)] + while pending: + path, candidate = pending.pop() + if not isinstance(candidate, Mapping): + continue + if isinstance(candidate.get("$ref"), str): + return _json_path_child(path, "$ref") + for key, child in candidate.items(): + if key in _DRAFT7_LITERAL_KEYWORDS: + continue + child_path = _json_path_child(path, key) + if isinstance(child, Mapping): + pending.append((child_path, child)) + elif isinstance(child, Sequence) and not isinstance(child, str): + pending.extend( + (_json_path_child(child_path, index), item) + for index, item in enumerate(child) + ) + return None @dataclass(frozen=True) @@ -1673,6 +1704,12 @@ def fetch_structured_output_schema( f"(stage=schema, path={error.json_path}): {error.message}" ) raise LLMNodeError(msg) from error + if ref_path := _find_draft7_ref_path(schema): + msg = ( + "Invalid structured output schema " + f"(stage=schema, path={ref_path}): $ref is not supported" + ) + raise LLMNodeError(msg) return schema @staticmethod diff --git a/tests/nodes/llm/test_node.py b/tests/nodes/llm/test_node.py index 1c82c425..d5457656 100644 --- a/tests/nodes/llm/test_node.py +++ b/tests/nodes/llm/test_node.py @@ -240,6 +240,52 @@ def test_fetch_structured_output_schema_checks_draft7_schema() -> None: assert "path=$.type" in str(exc_info.value) +@pytest.mark.parametrize( + ("schema", "path"), + [ + ( + { + "definitions": {"value": {"type": "string"}}, + "$ref": "#/definitions/value", + }, + "$['$ref']", + ), + ( + {"allOf": [{"$ref": "http://127.0.0.1/schema"}]}, + "$.allOf[0]['$ref']", + ), + ], + ids=["local", "nested-remote"], +) +def test_fetch_structured_output_schema_rejects_refs( + schema: dict[str, Any], + path: str, +) -> None: + with pytest.raises(LLMNodeError) as exc_info: + LLMNode.fetch_structured_output_schema( + structured_output={"schema": schema}, + ) + + assert "stage=schema" in str(exc_info.value) + assert f"path={path}" in str(exc_info.value) + assert "$ref is not supported" in str(exc_info.value) + + +def test_fetch_structured_output_schema_allows_ref_as_output_field_name() -> None: + schema = { + "type": "object", + "properties": {"$ref": {"type": "string"}}, + "enum": [{"$ref": "ordinary output value"}], + } + + assert ( + LLMNode.fetch_structured_output_schema( + structured_output={"schema": schema}, + ) + == schema + ) + + _RESULT_SCHEMA = { "type": "object", "properties": { From a5d84cfdab14cc41cb92985db444cfe14391ee73 Mon Sep 17 00:00:00 2001 From: -LAN- Date: Fri, 7 Aug 2026 09:28:08 +0800 Subject: [PATCH 3/3] fix(llm): disable remote schema retrieval Use an empty referencing registry so external schema references fail without network access while local references continue to work. --- pyproject.toml | 2 +- src/graphon/nodes/llm/node.py | 51 ++++++++---------------------- tests/nodes/llm/test_node.py | 59 +++++++++++++++-------------------- uv.lock | 2 +- 4 files changed, 40 insertions(+), 74 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 613538cf..b360aa92 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,7 +13,7 @@ dependencies = [ 'defusedxml>=0.7', 'httpx>=0.28', 'json-repair>=0.55', - 'jsonschema>=4', + 'jsonschema>=4.18', 'odfdo>=3.22.8', 'orjson>=3', 'pandas[excel]>=2.1', diff --git a/src/graphon/nodes/llm/node.py b/src/graphon/nodes/llm/node.py index 3bec6c19..e69f1578 100644 --- a/src/graphon/nodes/llm/node.py +++ b/src/graphon/nodes/llm/node.py @@ -12,6 +12,8 @@ from typing import Any, Literal, assert_never, override from jsonschema import Draft7Validator, SchemaError, ValidationError +from referencing import Registry +from referencing.exceptions import Unresolvable from graphon.entities.graph_init_params import GraphInitParams from graphon.enums import ( @@ -104,37 +106,7 @@ PromptMessageContentType.AUDIO: FileType.AUDIO, PromptMessageContentType.DOCUMENT: FileType.DOCUMENT, } -_DRAFT7_LITERAL_KEYWORDS = frozenset({"const", "default", "enum", "examples"}) - - -def _json_path_child(path: str, key: str | int) -> str: - if isinstance(key, int): - return f"{path}[{key}]" - if key.isidentifier(): - return f"{path}.{key}" - return f"{path}[{key!r}]" - - -def _find_draft7_ref_path(schema: Mapping[str, Any]) -> str | None: - pending: list[tuple[str, Any]] = [("$", schema)] - while pending: - path, candidate = pending.pop() - if not isinstance(candidate, Mapping): - continue - if isinstance(candidate.get("$ref"), str): - return _json_path_child(path, "$ref") - for key, child in candidate.items(): - if key in _DRAFT7_LITERAL_KEYWORDS: - continue - child_path = _json_path_child(path, key) - if isinstance(child, Mapping): - pending.append((child_path, child)) - elif isinstance(child, Sequence) and not isinstance(child, str): - pending.extend( - (_json_path_child(child_path, index), item) - for index, item in enumerate(child) - ) - return None +_NO_REMOTE_SCHEMA_REGISTRY = Registry() @dataclass(frozen=True) @@ -1129,13 +1101,22 @@ def _validate_structured_output_result( ) raise LLMNodeError(msg) try: - Draft7Validator(json_schema).validate(structured_output) + Draft7Validator( + json_schema, + registry=_NO_REMOTE_SCHEMA_REGISTRY, + ).validate(structured_output) except ValidationError as error: msg = ( "Structured output validation failed " f"(stage=result, path={error.json_path}): {error.message}" ) raise LLMNodeError(msg) from error + except Unresolvable as error: + msg = ( + "Structured output validation failed " + "(stage=result, path=$): schema reference could not be resolved" + ) + raise LLMNodeError(msg) from error @staticmethod def _finalize_streaming_usage( @@ -1704,12 +1685,6 @@ def fetch_structured_output_schema( f"(stage=schema, path={error.json_path}): {error.message}" ) raise LLMNodeError(msg) from error - if ref_path := _find_draft7_ref_path(schema): - msg = ( - "Invalid structured output schema " - f"(stage=schema, path={ref_path}): $ref is not supported" - ) - raise LLMNodeError(msg) return schema @staticmethod diff --git a/tests/nodes/llm/test_node.py b/tests/nodes/llm/test_node.py index d5457656..11aca21d 100644 --- a/tests/nodes/llm/test_node.py +++ b/tests/nodes/llm/test_node.py @@ -240,49 +240,40 @@ def test_fetch_structured_output_schema_checks_draft7_schema() -> None: assert "path=$.type" in str(exc_info.value) -@pytest.mark.parametrize( - ("schema", "path"), - [ - ( - { - "definitions": {"value": {"type": "string"}}, - "$ref": "#/definitions/value", - }, - "$['$ref']", - ), - ( - {"allOf": [{"$ref": "http://127.0.0.1/schema"}]}, - "$.allOf[0]['$ref']", - ), - ], - ids=["local", "nested-remote"], -) -def test_fetch_structured_output_schema_rejects_refs( - schema: dict[str, Any], - path: str, +def test_final_structured_output_validation_disables_remote_ref_retrieval( + monkeypatch: pytest.MonkeyPatch, ) -> None: + def fail_urlopen(*_args: Any, **_kwargs: Any) -> None: + pytest.fail("remote schema retrieval was attempted") + + monkeypatch.setattr("urllib.request.urlopen", fail_urlopen) + with pytest.raises(LLMNodeError) as exc_info: - LLMNode.fetch_structured_output_schema( - structured_output={"schema": schema}, + LLMNode._validate_structured_output_result( + structured_output={}, + json_schema={"$ref": "http://127.0.0.1/schema"}, ) - assert "stage=schema" in str(exc_info.value) - assert f"path={path}" in str(exc_info.value) - assert "$ref is not supported" in str(exc_info.value) + assert "stage=result" in str(exc_info.value) + assert "path=$" in str(exc_info.value) + assert "schema reference could not be resolved" in str(exc_info.value) -def test_fetch_structured_output_schema_allows_ref_as_output_field_name() -> None: +def test_final_structured_output_validation_supports_local_refs() -> None: schema = { - "type": "object", - "properties": {"$ref": {"type": "string"}}, - "enum": [{"$ref": "ordinary output value"}], + "definitions": { + "result": { + "type": "object", + "properties": {"ok": {"type": "boolean"}}, + "required": ["ok"], + }, + }, + "$ref": "#/definitions/result", } - assert ( - LLMNode.fetch_structured_output_schema( - structured_output={"schema": schema}, - ) - == schema + LLMNode._validate_structured_output_result( + structured_output={"ok": True}, + json_schema=schema, ) diff --git a/uv.lock b/uv.lock index 4896f3c7..a40a8d68 100644 --- a/uv.lock +++ b/uv.lock @@ -366,7 +366,7 @@ requires-dist = [ { name = "defusedxml", specifier = ">=0.7" }, { name = "httpx", specifier = ">=0.28" }, { name = "json-repair", specifier = ">=0.55" }, - { name = "jsonschema", specifier = ">=4" }, + { name = "jsonschema", specifier = ">=4.18" }, { name = "odfdo", specifier = ">=3.22.8" }, { name = "orjson", specifier = ">=3" }, { name = "pandas", extras = ["excel"], specifier = ">=2.1" },