From ea37904c756bd4a02e75013b959a12b27742dada Mon Sep 17 00:00:00 2001 From: Stream Date: Wed, 21 Jan 2026 19:30:46 +0800 Subject: [PATCH 1/2] refactor: unify structured output with pydantic model Signed-off-by: Stream --- api/core/llm_generator/llm_generator.py | 101 +++++------------- api/core/llm_generator/output_models.py | 34 ++++++ .../output_parser/structured_output.py | 89 ++++++++++++++- .../test_structured_output_parser.py | 71 +++++++++++- 4 files changed, 216 insertions(+), 79 deletions(-) create mode 100644 api/core/llm_generator/output_models.py diff --git a/api/core/llm_generator/llm_generator.py b/api/core/llm_generator/llm_generator.py index fd769f6a83..d29332c3fa 100644 --- a/api/core/llm_generator/llm_generator.py +++ b/api/core/llm_generator/llm_generator.py @@ -6,6 +6,11 @@ from typing import Any, Protocol, cast import json_repair +from core.llm_generator.output_models import ( + CodeNodeStructuredOutput, + InstructionModifyOutput, + SuggestedQuestionsOutput, +) from core.llm_generator.output_parser.rule_config_generator import RuleConfigGeneratorOutputParser from core.llm_generator.output_parser.suggested_questions_after_answer import SuggestedQuestionsAfterAnswerOutputParser from core.llm_generator.prompts import ( @@ -470,7 +475,7 @@ class LLMGenerator: *prompt_messages, ] - from core.llm_generator.output_parser.structured_output import invoke_llm_with_structured_output + from core.llm_generator.output_parser.structured_output import invoke_llm_with_pydantic_model # Get model instance and schema provider = model_config.get("provider", "") @@ -487,15 +492,13 @@ class LLMGenerator: return cls._error_response(f"Model schema not found for {model_name}") model_parameters = model_config.get("completion_params", {}) - json_schema = cls._get_code_node_json_schema() - try: - response = invoke_llm_with_structured_output( + response = invoke_llm_with_pydantic_model( provider=provider, model_schema=model_schema, model_instance=model_instance, prompt_messages=complete_messages, - json_schema=json_schema, + output_model=CodeNodeStructuredOutput, model_parameters=model_parameters, stream=False, tenant_id=tenant_id, @@ -541,7 +544,7 @@ class LLMGenerator: from sqlalchemy import select from sqlalchemy.orm import Session - from core.llm_generator.output_parser.structured_output import invoke_llm_with_structured_output + from core.llm_generator.output_parser.structured_output import invoke_llm_with_pydantic_model from services.workflow_service import WorkflowService # Get workflow context (reuse existing logic) @@ -602,15 +605,13 @@ class LLMGenerator: completion_params = model_config.get("completion_params", {}) if model_config else {} model_parameters = {**completion_params, "max_tokens": 256} - json_schema = cls._get_suggested_questions_json_schema() - try: - response = invoke_llm_with_structured_output( + response = invoke_llm_with_pydantic_model( provider=model_instance.provider, model_schema=model_schema, model_instance=model_instance, prompt_messages=prompt_messages, - json_schema=json_schema, + output_model=SuggestedQuestionsOutput, model_parameters=model_parameters, stream=False, tenant_id=tenant_id, @@ -644,58 +645,6 @@ Sources: {", ".join(sources)} Target: {parameter_info.get("name")}({param_type}) - {param_desc} Output 3 short, practical questions in {language}.""" - @classmethod - def _get_suggested_questions_json_schema(cls) -> dict: - """Return JSON Schema for suggested questions.""" - return { - "type": "object", - "properties": { - "questions": { - "type": "array", - "items": {"type": "string"}, - "minItems": 3, - "maxItems": 3, - "description": "3 suggested questions", - }, - }, - "required": ["questions"], - } - - @classmethod - def _get_code_node_json_schema(cls) -> dict: - """Return JSON Schema for structured output.""" - return { - "type": "object", - "properties": { - "variables": { - "type": "array", - "items": { - "type": "object", - "properties": { - "variable": {"type": "string", "description": "Variable name in code"}, - "value_selector": { - "type": "array", - "items": {"type": "string"}, - "description": "Path like [node_id, output_name]", - }, - }, - "required": ["variable", "value_selector"], - }, - }, - "code": {"type": "string", "description": "Generated code with main function"}, - "outputs": { - "type": "object", - "additionalProperties": { - "type": "object", - "properties": {"type": {"type": "string"}}, - }, - "description": "Output definitions, key is output name", - }, - "explanation": {"type": "string", "description": "Brief explanation of the code"}, - }, - "required": ["variables", "code", "outputs", "explanation"], - } - @classmethod def _get_upstream_nodes(cls, graph_dict: Mapping[str, Any], node_id: str) -> list[dict]: """ @@ -1011,6 +960,10 @@ Parameter: {parameter_info.get("name")} ({param_type}) - {parameter_info.get("de provider=model_config.get("provider", ""), model=model_config.get("name", ""), ) + model_name = model_config.get("name", "") + model_schema = model_instance.model_type_instance.get_model_schema(model_name, model_instance.credentials) + if not model_schema: + return {"error": f"Model schema not found for {model_name}"} match node_type: case "llm" | "agent": system_prompt = LLM_MODIFY_PROMPT_SYSTEM @@ -1034,20 +987,18 @@ Parameter: {parameter_info.get("name")} ({param_type}) - {parameter_info.get("de model_parameters = {"temperature": 0.4} try: - response: LLMResult = model_instance.invoke_llm( - prompt_messages=list(prompt_messages), model_parameters=model_parameters, stream=False - ) + from core.llm_generator.output_parser.structured_output import invoke_llm_with_pydantic_model - generated_raw = response.message.get_text_content() - first_brace = generated_raw.find("{") - last_brace = generated_raw.rfind("}") - if first_brace == -1 or last_brace == -1 or last_brace < first_brace: - raise ValueError(f"Could not find a valid JSON object in response: {generated_raw}") - json_str = generated_raw[first_brace : last_brace + 1] - data = json_repair.loads(json_str) - if not isinstance(data, dict): - raise TypeError(f"Expected a JSON object, but got {type(data).__name__}") - return data + response = invoke_llm_with_pydantic_model( + provider=model_instance.provider, + model_schema=model_schema, + model_instance=model_instance, + prompt_messages=list(prompt_messages), + output_model=InstructionModifyOutput, + model_parameters=model_parameters, + stream=False, + ) + return response.structured_output or {} except InvokeError as e: error = str(e) return {"error": f"Failed to generate code. Error: {error}"} diff --git a/api/core/llm_generator/output_models.py b/api/core/llm_generator/output_models.py new file mode 100644 index 0000000000..22c4f5e411 --- /dev/null +++ b/api/core/llm_generator/output_models.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +from pydantic import BaseModel, ConfigDict, Field + +from core.variables.types import SegmentType +from core.workflow.nodes.base.entities import VariableSelector + + +class SuggestedQuestionsOutput(BaseModel): + model_config = ConfigDict(extra="forbid") + + questions: list[str] = Field(min_length=3, max_length=3) + + +class CodeNodeOutput(BaseModel): + model_config = ConfigDict(extra="forbid") + + type: SegmentType + + +class CodeNodeStructuredOutput(BaseModel): + model_config = ConfigDict(extra="forbid") + + variables: list[VariableSelector] + code: str + outputs: dict[str, CodeNodeOutput] + explanation: str + + +class InstructionModifyOutput(BaseModel): + model_config = ConfigDict(extra="forbid") + + modified: str + message: str diff --git a/api/core/llm_generator/output_parser/structured_output.py b/api/core/llm_generator/output_parser/structured_output.py index 250acf14fd..7e931fed32 100644 --- a/api/core/llm_generator/output_parser/structured_output.py +++ b/api/core/llm_generator/output_parser/structured_output.py @@ -2,10 +2,10 @@ import json from collections.abc import Generator, Mapping, Sequence from copy import deepcopy from enum import StrEnum -from typing import Any, Literal, cast, overload +from typing import Any, Literal, TypeVar, cast, overload import json_repair -from pydantic import TypeAdapter, ValidationError +from pydantic import BaseModel, TypeAdapter, ValidationError from core.llm_generator.output_parser.errors import OutputParserError from core.llm_generator.output_parser.file_ref import convert_file_refs_in_output @@ -44,6 +44,9 @@ class SpecialModelType(StrEnum): OLLAMA = "ollama" +T = TypeVar("T", bound=BaseModel) + + @overload def invoke_llm_with_structured_output( *, @@ -129,7 +132,6 @@ def invoke_llm_with_structured_output( file IDs in the output will be automatically converted to File objects. :return: full response or stream response chunk generator result """ - # handle native json schema model_parameters_with_json_schema: dict[str, Any] = { **(model_parameters or {}), @@ -234,6 +236,87 @@ def invoke_llm_with_structured_output( return generator() +@overload +def invoke_llm_with_pydantic_model( + *, + provider: str, + model_schema: AIModelEntity, + model_instance: ModelInstance, + prompt_messages: Sequence[PromptMessage], + output_model: type[T], + model_parameters: Mapping | None = None, + tools: Sequence[PromptMessageTool] | None = None, + stop: list[str] | None = None, + stream: Literal[False] = False, + user: str | None = None, + callbacks: list[Callback] | None = None, + tenant_id: str | None = None, +) -> LLMResultWithStructuredOutput: ... + + +def invoke_llm_with_pydantic_model( + *, + provider: str, + model_schema: AIModelEntity, + model_instance: ModelInstance, + prompt_messages: Sequence[PromptMessage], + output_model: type[T], + model_parameters: Mapping | None = None, + tools: Sequence[PromptMessageTool] | None = None, + stop: list[str] | None = None, + stream: bool = False, + user: str | None = None, + callbacks: list[Callback] | None = None, + tenant_id: str | None = None, +) -> LLMResultWithStructuredOutput: + """ + Invoke large language model with a Pydantic output model. + + This helper generates a JSON schema from the Pydantic model, invokes the + structured-output LLM path, and validates the result in non-streaming mode. + """ + if stream: + raise ValueError("invoke_llm_with_pydantic_model only supports stream=False") + + json_schema = _schema_from_pydantic(output_model) + result = invoke_llm_with_structured_output( + provider=provider, + model_schema=model_schema, + model_instance=model_instance, + prompt_messages=prompt_messages, + json_schema=json_schema, + model_parameters=model_parameters, + tools=tools, + stop=stop, + stream=False, + user=user, + callbacks=callbacks, + tenant_id=tenant_id, + ) + + structured_output = result.structured_output + if structured_output is None: + raise OutputParserError("Structured output is empty") + + validated_output = _validate_structured_output(output_model, structured_output) + return result.model_copy(update={"structured_output": validated_output}) + + +def _schema_from_pydantic(output_model: type[BaseModel]) -> dict[str, Any]: + return output_model.model_json_schema() + + +def _validate_structured_output( + output_model: type[T], + structured_output: Mapping[str, Any], +) -> dict[str, Any]: + try: + validated_output = output_model.model_validate(structured_output) + except ValidationError as exc: + raise OutputParserError(f"Structured output validation failed: {exc}") from exc + return validated_output.model_dump(mode="python") + + def _handle_native_json_schema( provider: str, model_schema: AIModelEntity, diff --git a/api/tests/unit_tests/utils/structured_output_parser/test_structured_output_parser.py b/api/tests/unit_tests/utils/structured_output_parser/test_structured_output_parser.py index 9046f785d2..9742590cd4 100644 --- a/api/tests/unit_tests/utils/structured_output_parser/test_structured_output_parser.py +++ b/api/tests/unit_tests/utils/structured_output_parser/test_structured_output_parser.py @@ -2,9 +2,13 @@ from decimal import Decimal from unittest.mock import MagicMock, patch import pytest +from pydantic import BaseModel, ConfigDict from core.llm_generator.output_parser.errors import OutputParserError -from core.llm_generator.output_parser.structured_output import invoke_llm_with_structured_output +from core.llm_generator.output_parser.structured_output import ( + invoke_llm_with_pydantic_model, + invoke_llm_with_structured_output, +) from core.model_runtime.entities.llm_entities import ( LLMResult, LLMResultChunk, @@ -461,3 +465,68 @@ def test_model_specific_schema_preparation(): # For Gemini, the schema should not have additionalProperties and boolean should be converted to string assert "json_schema" in call_args.kwargs["model_parameters"] + + +class ExampleOutput(BaseModel): + model_config = ConfigDict(extra="forbid") + + name: str + + +def test_structured_output_with_pydantic_model(): + model_schema = get_model_entity("openai", "gpt-4o", support_structure_output=True) + model_instance = get_model_instance() + model_instance.invoke_llm.return_value = LLMResult( + model="gpt-4o", + message=AssistantPromptMessage(content='{"name": "test"}'), + usage=create_mock_usage(prompt_tokens=8, completion_tokens=4), + ) + + prompt_messages = [UserPromptMessage(content="Return a JSON object with name.")] + + result = invoke_llm_with_pydantic_model( + provider="openai", + model_schema=model_schema, + model_instance=model_instance, + prompt_messages=prompt_messages, + output_model=ExampleOutput, + stream=False, + ) + + assert isinstance(result, LLMResultWithStructuredOutput) + assert result.structured_output == {"name": "test"} + + +def test_structured_output_with_pydantic_model_streaming_rejected(): + model_schema = get_model_entity("openai", "gpt-4o", support_structure_output=True) + model_instance = get_model_instance() + + with pytest.raises(ValueError): + invoke_llm_with_pydantic_model( + provider="openai", + model_schema=model_schema, + model_instance=model_instance, + prompt_messages=[UserPromptMessage(content="test")], + output_model=ExampleOutput, + stream=True, + ) + + +def test_structured_output_with_pydantic_model_validation_error(): + model_schema = get_model_entity("openai", "gpt-4o", support_structure_output=True) + model_instance = get_model_instance() + model_instance.invoke_llm.return_value = LLMResult( + model="gpt-4o", + message=AssistantPromptMessage(content='{"name": 123}'), + usage=create_mock_usage(prompt_tokens=8, completion_tokens=4), + ) + + with pytest.raises(OutputParserError): + invoke_llm_with_pydantic_model( + provider="openai", + model_schema=model_schema, + model_instance=model_instance, + prompt_messages=[UserPromptMessage(content="test")], + output_model=ExampleOutput, + stream=False, + ) From 39ec2b3277083ece2c36df4eeae5ecf6c8f1f9b4 Mon Sep 17 00:00:00 2001 From: zhsama Date: Tue, 27 Jan 2026 19:39:32 +0800 Subject: [PATCH 2/2] feat: Add file type support to LLM node JSON schema editor --- .../nodes/_base/components/variable/utils.ts | 47 ++++++++++++++----- .../visual-editor/edit-card/index.tsx | 4 ++ .../visual-editor/hooks.ts | 38 ++++++++++++++- .../components/workflow/nodes/llm/types.ts | 6 ++- .../components/workflow/nodes/llm/utils.ts | 9 ++-- 5 files changed, 87 insertions(+), 17 deletions(-) diff --git a/web/app/components/workflow/nodes/_base/components/variable/utils.ts b/web/app/components/workflow/nodes/_base/components/variable/utils.ts index 2f83945dc2..c3a3201bcd 100644 --- a/web/app/components/workflow/nodes/_base/components/variable/utils.ts +++ b/web/app/components/workflow/nodes/_base/components/variable/utils.ts @@ -60,10 +60,7 @@ import { import { VAR_REGEX } from '@/config' import { AppModeEnum } from '@/types/app' import { OUTPUT_FILE_SUB_VARIABLES } from '../../../constants' -import { - - Type, -} from '../../../llm/types' +import { FILE_REF_FORMAT, Type } from '../../../llm/types' import { VarType as ToolVarType } from '../../../tool/types' export const isSystemVar = (valueSelector: ValueSelector) => { @@ -122,7 +119,16 @@ export const inputVarTypeToVarType = (type: InputVarType): VarType => { ) } -const structTypeToVarType = (type: Type, isArray?: boolean): VarType => { +const structTypeToVarType = ( + type: Type, + isArray?: boolean, + format?: string, + itemsFormat?: string, +): VarType => { + if (isArray && itemsFormat === FILE_REF_FORMAT) + return VarType.arrayFile + if (!isArray && format === FILE_REF_FORMAT) + return VarType.file if (isArray) { return ( ( @@ -175,6 +181,7 @@ const findExceptVarInStructuredProperties = ( const isObj = item.type === Type.object const isArray = item.type === Type.array const arrayType = item.items?.type + const arrayFormat = item.items?.format if ( !isObj @@ -184,6 +191,8 @@ const findExceptVarInStructuredProperties = ( type: structTypeToVarType( isArray ? arrayType! : item.type, isArray, + item.format, + arrayFormat, ), }, [key], @@ -215,6 +224,7 @@ const findExceptVarInStructuredOutput = ( const isObj = item.type === Type.object const isArray = item.type === Type.array const arrayType = item.items?.type + const arrayFormat = item.items?.format if ( !isObj && !filterVar( @@ -223,6 +233,8 @@ const findExceptVarInStructuredOutput = ( type: structTypeToVarType( isArray ? arrayType! : item.type, isArray, + item.format, + arrayFormat, ), }, [key], @@ -1144,8 +1156,14 @@ export const getVarType = ({ return currProperties = currProperties.properties[key] - if (isLast) - type = structTypeToVarType(currProperties?.type) + if (isLast) { + if (currProperties?.format === FILE_REF_FORMAT) + type = VarType.file + else if (currProperties?.type === Type.array && currProperties?.items?.format === FILE_REF_FORMAT) + type = VarType.arrayFile + else + type = structTypeToVarType(currProperties?.type) + } }) return type } @@ -1946,15 +1964,20 @@ const varToValueSelectorList = ( Object.keys( (v.children as StructuredOutput)?.schema?.properties || {}, ).forEach((key) => { - const type = (v.children as StructuredOutput)?.schema?.properties[key].type + const schemaProperty = (v.children as StructuredOutput)?.schema?.properties[key] + const type = schemaProperty?.type const isArray = type === Type.array - const arrayType = (v.children as StructuredOutput)?.schema?.properties[ - key - ].items?.type + const arrayType = schemaProperty?.items?.type + const arrayFormat = schemaProperty?.items?.format varToValueSelectorList( { variable: key, - type: structTypeToVarType(isArray ? arrayType! : type, isArray), + type: structTypeToVarType( + isArray ? arrayType! : type, + isArray, + schemaProperty?.format, + arrayFormat, + ), }, [...parentValueSelector, v.variable], res, diff --git a/web/app/components/workflow/nodes/llm/components/json-schema-config-modal/visual-editor/edit-card/index.tsx b/web/app/components/workflow/nodes/llm/components/json-schema-config-modal/visual-editor/edit-card/index.tsx index 3cf9003051..b3713ae60a 100644 --- a/web/app/components/workflow/nodes/llm/components/json-schema-config-modal/visual-editor/edit-card/index.tsx +++ b/web/app/components/workflow/nodes/llm/components/json-schema-config-modal/visual-editor/edit-card/index.tsx @@ -44,17 +44,21 @@ const TYPE_OPTIONS = [ { value: Type.number, text: 'number' }, { value: Type.boolean, text: 'boolean' }, { value: Type.object, text: 'object' }, + { value: Type.file, text: 'file' }, { value: ArrayType.string, text: 'array[string]' }, { value: ArrayType.number, text: 'array[number]' }, { value: ArrayType.object, text: 'array[object]' }, + { value: ArrayType.file, text: 'array[file]' }, ] const MAXIMUM_DEPTH_TYPE_OPTIONS = [ { value: Type.string, text: 'string' }, { value: Type.number, text: 'number' }, { value: Type.boolean, text: 'boolean' }, + { value: Type.file, text: 'file' }, { value: ArrayType.string, text: 'array[string]' }, { value: ArrayType.number, text: 'array[number]' }, + { value: ArrayType.file, text: 'array[file]' }, ] const EditCard: FC = ({ diff --git a/web/app/components/workflow/nodes/llm/components/json-schema-config-modal/visual-editor/hooks.ts b/web/app/components/workflow/nodes/llm/components/json-schema-config-modal/visual-editor/hooks.ts index 6159028c21..a2b2d8e2b8 100644 --- a/web/app/components/workflow/nodes/llm/components/json-schema-config-modal/visual-editor/hooks.ts +++ b/web/app/components/workflow/nodes/llm/components/json-schema-config-modal/visual-editor/hooks.ts @@ -4,7 +4,7 @@ import type { EditData } from './edit-card' import { noop } from 'es-toolkit/function' import { produce } from 'immer' import Toast from '@/app/components/base/toast' -import { ArrayType, Type } from '../../../types' +import { ArrayType, FILE_REF_FORMAT, Type } from '../../../types' import { findPropertyWithPath } from '../../../utils' import { useMittContext } from './context' import { useVisualEditorStore } from './store' @@ -133,6 +133,7 @@ export const useSchemaNodeOperations = (props: VisualEditorProps) => { } if (schema.type === Type.array) delete schema.items + delete schema.format switch (newType) { case Type.object: schema.type = Type.object @@ -140,6 +141,10 @@ export const useSchemaNodeOperations = (props: VisualEditorProps) => { schema.required = [] schema.additionalProperties = false break + case Type.file: + schema.type = Type.string + schema.format = FILE_REF_FORMAT + break case ArrayType.string: schema.type = Type.array schema.items = { @@ -167,6 +172,13 @@ export const useSchemaNodeOperations = (props: VisualEditorProps) => { additionalProperties: false, } break + case ArrayType.file: + schema.type = Type.array + schema.items = { + type: Type.string, + format: FILE_REF_FORMAT, + } + break default: schema.type = newType as Type } @@ -309,6 +321,7 @@ export const useSchemaNodeOperations = (props: VisualEditorProps) => { } if (schema.type === Type.array) delete schema.items + delete schema.format switch (newType) { case Type.object: schema.type = Type.object @@ -316,6 +329,10 @@ export const useSchemaNodeOperations = (props: VisualEditorProps) => { schema.required = [] schema.additionalProperties = false break + case Type.file: + schema.type = Type.string + schema.format = FILE_REF_FORMAT + break case ArrayType.string: schema.type = Type.array schema.items = { @@ -343,6 +360,13 @@ export const useSchemaNodeOperations = (props: VisualEditorProps) => { additionalProperties: false, } break + case ArrayType.file: + schema.type = Type.array + schema.items = { + type: Type.string, + format: FILE_REF_FORMAT, + } + break default: schema.type = newType as Type } @@ -398,6 +422,7 @@ export const useSchemaNodeOperations = (props: VisualEditorProps) => { } if (schema.type === Type.array) delete schema.items + delete schema.format switch (newType) { case Type.object: schema.type = Type.object @@ -405,6 +430,10 @@ export const useSchemaNodeOperations = (props: VisualEditorProps) => { schema.required = [] schema.additionalProperties = false break + case Type.file: + schema.type = Type.string + schema.format = FILE_REF_FORMAT + break case ArrayType.string: schema.type = Type.array schema.items = { @@ -432,6 +461,13 @@ export const useSchemaNodeOperations = (props: VisualEditorProps) => { additionalProperties: false, } break + case ArrayType.file: + schema.type = Type.array + schema.items = { + type: Type.string, + format: FILE_REF_FORMAT, + } + break default: schema.type = newType as Type } diff --git a/web/app/components/workflow/nodes/llm/types.ts b/web/app/components/workflow/nodes/llm/types.ts index 5b15f83ac6..9738b74ea2 100644 --- a/web/app/components/workflow/nodes/llm/types.ts +++ b/web/app/components/workflow/nodes/llm/types.ts @@ -20,6 +20,8 @@ export type LLMNodeType = CommonNodeType & { reasoning_format?: 'tagged' | 'separated' } +export const FILE_REF_FORMAT = 'dify-file-ref' + export enum Type { string = 'string', number = 'number', @@ -38,12 +40,13 @@ export enum ArrayType { number = 'array[number]', boolean = 'array[boolean]', object = 'array[object]', + file = 'array[file]', } export type TypeWithArray = Type | ArrayType type ArrayItemType = Exclude -export type ArrayItems = Omit & { type: ArrayItemType } +export type ArrayItems = Omit & { type: ArrayItemType; format?: string } export type SchemaEnumType = string[] | number[] @@ -54,6 +57,7 @@ export type Field = { } required?: string[] // Key of required properties in object description?: string + format?: string items?: ArrayItems // Array has items. Define the item type enum?: SchemaEnumType // Enum values additionalProperties?: false // Required in object by api. Just set false diff --git a/web/app/components/workflow/nodes/llm/utils.ts b/web/app/components/workflow/nodes/llm/utils.ts index 604a1f8408..9922529cb5 100644 --- a/web/app/components/workflow/nodes/llm/utils.ts +++ b/web/app/components/workflow/nodes/llm/utils.ts @@ -2,21 +2,24 @@ import type { ValidationError } from 'jsonschema' import type { ArrayItems, Field, LLMNodeType } from './types' import { z } from 'zod' import { draft07Validator, forbidBooleanProperties } from '@/utils/validators' -import { ArrayType, Type } from './types' +import { ArrayType, FILE_REF_FORMAT, Type } from './types' export const checkNodeValid = (_payload: LLMNodeType) => { return true } export const getFieldType = (field: Field) => { - const { type, items, enum: enums } = field + const { type, items, enum: enums, format } = field + if (format === FILE_REF_FORMAT) + return Type.file if (field.schemaType === 'file') return Type.file if (enums && enums.length > 0) return Type.enumType if (type !== Type.array || !items) return type - + if (items.format === FILE_REF_FORMAT || items.type === Type.file) + return ArrayType.file return ArrayType[items.type as keyof typeof ArrayType] }