chore: fix mypy issue

This commit is contained in:
Novice 2025-04-11 16:36:31 +08:00
parent ae62d089db
commit 86eefb7ed8

View File

@ -202,14 +202,18 @@ class LLMNode(BaseNode[LLMNodeData]):
break break
outputs = {"text": result_text, "usage": jsonable_encoder(usage), "finish_reason": finish_reason} outputs = {"text": result_text, "usage": jsonable_encoder(usage), "finish_reason": finish_reason}
if self.node_data.structured_output_enabled and self.node_data.structured_output: if self.node_data.structured_output_enabled and self.node_data.structured_output:
structured_output = {} structured_output: dict[str, Any] | list[Any] = {}
try: try:
structured_output = json.loads(result_text) parsed = json.loads(result_text)
if not isinstance(parsed, (dict, list)):
raise LLMNodeError(f"Failed to parse structured output: {result_text}")
structured_output = parsed
except json.JSONDecodeError as e: except json.JSONDecodeError as e:
# if the result_text is not a valid json, try to repair it # if the result_text is not a valid json, try to repair it
structured_output = json_repair.loads(result_text) parsed = json_repair.loads(result_text)
if not isinstance(structured_output, dict | list): if not isinstance(parsed, (dict, list)):
raise LLMNodeError(f"Failed to parse structured output: {result_text}") raise LLMNodeError(f"Failed to parse structured output: {result_text}")
structured_output = parsed
outputs["structured_output"] = structured_output outputs["structured_output"] = structured_output
yield RunCompletedEvent( yield RunCompletedEvent(
run_result=NodeRunResult( run_result=NodeRunResult(
@ -1066,17 +1070,13 @@ class LLMNode(BaseNode[LLMNodeData]):
model_schema = model_type_instance.get_model_schema(model_name, model_credentials) model_schema = model_type_instance.get_model_schema(model_name, model_credentials)
return model_schema return model_schema
def _fetch_structured_output_schema(self) -> dict: def _fetch_structured_output_schema(self) -> dict[str, Any]:
""" """
Fetch the structured output schema for language models. Fetch the structured output schema from the node data.
This method retrieves and validates the JSON schema defined in the node data Returns:
for structured output formatting. It ensures the schema is properly formatted dict[str, Any]: The structured output schema
and can be used for structured response generation.
:return: The validated JSON schema as a dictionary
""" """
# Extract and validate schema
if not self.node_data.structured_output: if not self.node_data.structured_output:
raise LLMNodeError("Please provide a valid structured output schema") raise LLMNodeError("Please provide a valid structured output schema")
structured_output_schema = json.dumps(self.node_data.structured_output.get("schema", {}), ensure_ascii=False) structured_output_schema = json.dumps(self.node_data.structured_output.get("schema", {}), ensure_ascii=False)
@ -1085,9 +1085,11 @@ class LLMNode(BaseNode[LLMNodeData]):
try: try:
schema = json.loads(structured_output_schema) schema = json.loads(structured_output_schema)
if not isinstance(schema, dict):
raise LLMNodeError("structured_output_schema must be a JSON object")
return schema
except json.JSONDecodeError: except json.JSONDecodeError:
raise LLMNodeError("structured_output_schema is not valid JSON format") raise LLMNodeError("structured_output_schema is not valid JSON format")
return schema
def _check_model_structured_output_support(self) -> Optional[bool]: def _check_model_structured_output_support(self) -> Optional[bool]:
""" """