Skip to content

Commit ce3a870

Browse files
committed
feat: Support parameter extraction nodes
1 parent cfef12b commit ce3a870

File tree

16 files changed

+827
-27
lines changed

16 files changed

+827
-27
lines changed

apps/application/flow/step_node/__init__.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
from .loop_node import *
2222
from .loop_start_node import *
2323
from .mcp_node import BaseMcpNode
24+
from .parameter_extraction_node import BaseParameterExtractionNode
2425
from .question_node import *
2526
from .reranker_node import *
2627
from .search_knowledge_node import *
@@ -43,7 +44,7 @@
4344
BaseVideoUnderstandNode,
4445
BaseIntentNode, BaseLoopNode, BaseLoopStartStepNode,
4546
BaseLoopContinueNode,
46-
BaseLoopBreakNode, BaseVariableSplittingNode]
47+
BaseLoopBreakNode, BaseVariableSplittingNode, BaseParameterExtractionNode]
4748

4849

4950
def get_node(node_type):
Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
# coding=utf-8
2+
"""
3+
@project: MaxKB
4+
@Author:虎虎
5+
@file: __init__.py.py
6+
@date:2025/10/13 14:56
7+
@desc:
8+
"""
9+
from .impl import *
Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
# coding=utf-8
2+
3+
from typing import Type
4+
5+
from django.utils.translation import gettext_lazy as _
6+
from rest_framework import serializers
7+
8+
from application.flow.i_step_node import INode, NodeResult
9+
10+
11+
class VariableSplittingNodeParamsSerializer(serializers.Serializer):
12+
input_variable = serializers.ListField(required=True,
13+
label=_("input variable"))
14+
15+
variable_list = serializers.ListField(required=True,
16+
label=_("Split variables"))
17+
18+
model_params_setting = serializers.DictField(required=False,
19+
label=_("Model parameter settings"))
20+
21+
model_id = serializers.CharField(required=True, label=_("Model id"))
22+
23+
24+
class IParameterExtractionNode(INode):
25+
type = 'parameter-extraction-node'
26+
27+
def get_node_params_serializer_class(self) -> Type[serializers.Serializer]:
28+
return VariableSplittingNodeParamsSerializer
29+
30+
def _run(self):
31+
input_variable = self.workflow_manage.get_reference_field(
32+
self.node_params_serializer.data.get('input_variable')[0],
33+
self.node_params_serializer.data.get('input_variable')[1:])
34+
return self.execute(input_variable, self.node_params_serializer.data['variable_list'],
35+
self.node_params_serializer.data['model_params_setting'],
36+
self.node_params_serializer.data['model_id'])
37+
38+
def execute(self, input_variable, variable_list, model_params_setting, model_id, **kwargs) -> NodeResult:
39+
pass
Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
# coding=utf-8
2+
"""
3+
@project: MaxKB
4+
@Author:虎虎
5+
@file: __init__.py.py
6+
@date:2025/10/13 15:01
7+
@desc:
8+
"""
9+
from .base_parameter_extraction_node import *
Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,112 @@
1+
# coding=utf-8
2+
"""
3+
@project: MaxKB
4+
@Author:虎虎
5+
@file: base_variable_splitting_node.py
6+
@date:2025/10/13 15:02
7+
@desc:
8+
"""
9+
import json
10+
import re
11+
12+
from django.db.models import QuerySet
13+
from langchain_core.messages import HumanMessage
14+
from langchain_core.prompts import PromptTemplate
15+
16+
from application.flow.i_step_node import NodeResult
17+
from application.flow.step_node.parameter_extraction_node.i_parameter_extraction_node import IParameterExtractionNode
18+
from models_provider.models import Model
19+
from models_provider.tools import get_model_instance_by_model_workspace_id, get_model_credential
20+
21+
prompt = """
22+
Please strictly process the text according to the following requirements:
23+
**Task**:
24+
Extract specified field information from given text
25+
26+
**Enter text**:
27+
{{question}}
28+
29+
**Extract configuration**:
30+
{{properties}}
31+
32+
**Rule**:
33+
- Strictly follow the data and field of Extract configuration
34+
- If not found, use null value
35+
- Only return pure JSON without additional text
36+
- Keep the string format neat
37+
"""
38+
39+
40+
def get_default_model_params_setting(model_id):
41+
model = QuerySet(Model).filter(id=model_id).first()
42+
credential = get_model_credential(model.provider, model.model_type, model.model_name)
43+
model_params_setting = credential.get_model_params_setting_form(
44+
model.model_name).get_default_form_data()
45+
return model_params_setting
46+
47+
48+
def generate_properties(variable_list):
49+
return {variable['field']: {'type': variable['parameter_type'], 'description': variable['desc'],
50+
'title': variable['label']} for variable in
51+
variable_list}
52+
53+
54+
def generate_example(variable_list):
55+
return {variable['field']: None for variable in variable_list}
56+
57+
58+
def generate_content(input_variable, variable_list):
59+
properties = generate_properties(variable_list)
60+
prompt_template = PromptTemplate.from_template(prompt, template_format='jinja2')
61+
value = prompt_template.format(properties=properties, question=input_variable)
62+
return value
63+
64+
65+
def json_loads(response, expected_fields):
66+
if not response or not isinstance(response, str):
67+
return {field: None for field in expected_fields}
68+
69+
cleaned = response.strip()
70+
71+
extraction_strategies = [
72+
lambda: json.loads(cleaned),
73+
lambda: json.loads(re.search(r'```(?:json)?\s*(\{.*?\})\s*```', cleaned, re.DOTALL).group(1)),
74+
lambda: json.loads(re.search(r'(\{[\s\S]*\})', cleaned).group(1)),
75+
]
76+
for strategy in extraction_strategies:
77+
try:
78+
result = strategy()
79+
return result
80+
except:
81+
continue
82+
return generate_example(expected_fields)
83+
84+
85+
class BaseParameterExtractionNode(IParameterExtractionNode):
86+
87+
def save_context(self, details, workflow_manage):
88+
for key, value in details.get('result').items():
89+
self.context['key'] = value
90+
self.context['result'] = details.get('result')
91+
92+
def execute(self, input_variable, variable_list, model_params_setting, model_id, **kwargs) -> NodeResult:
93+
if model_params_setting is None:
94+
model_params_setting = get_default_model_params_setting(model_id)
95+
workspace_id = self.workflow_manage.get_body().get('workspace_id')
96+
chat_model = get_model_instance_by_model_workspace_id(model_id, workspace_id,
97+
**model_params_setting)
98+
content = generate_content(input_variable, variable_list)
99+
response = chat_model.invoke([HumanMessage(content=content)])
100+
result = json_loads(response.content, variable_list)
101+
return NodeResult({'result': result, **result}, {})
102+
103+
def get_details(self, index: int, **kwargs):
104+
return {
105+
'name': self.node.properties.get('stepName'),
106+
"index": index,
107+
'run_time': self.context.get('run_time'),
108+
'type': self.node.type,
109+
'result': self.context.get('result'),
110+
'status': self.status,
111+
'err_message': self.err_message
112+
}
Lines changed: 11 additions & 0 deletions
Loading

ui/src/enums/application.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ export enum WorkflowType {
3434
LoopBreakNode = 'loop-break-node',
3535
VariableSplittingNode = 'variable-splitting-node',
3636
VideoUnderstandNode = 'video-understand-node',
37+
ParameterExtractionNode = 'parameter-extraction-node',
3738
}
3839
export enum WorkflowMode {
3940
// 应用工作流

ui/src/locales/lang/en-US/views/application-workflow.ts

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -407,6 +407,20 @@ You are a master of problem optimization, adept at accurately inferring user int
407407
placeholder: 'Please enter expression',
408408
},
409409
},
410+
parameterExtractionNode: {
411+
label: '參數提取',
412+
text: '利用 AI 模型提取結構化參數',
413+
result: '結果',
414+
selectVariables: {
415+
label: '選擇變數',
416+
placeholder: '請選擇變數',
417+
},
418+
extractParameters: {
419+
label: '提取參數',
420+
desc: '描述',
421+
parameterType: '參數類型',
422+
},
423+
},
410424
},
411425
compare: {
412426
is_null: 'Is null',

ui/src/locales/lang/zh-CN/views/application-workflow.ts

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -420,6 +420,20 @@ export default {
420420
placeholder: '请输入表达式',
421421
},
422422
},
423+
parameterExtractionNode: {
424+
label: '参数提取',
425+
text: '利用 AI 模型提取结构化参数',
426+
result: '结果',
427+
selectVariables: {
428+
label: '选择变量',
429+
placeholder: '请选择变量',
430+
},
431+
extractParameters: {
432+
label: '提取参数',
433+
desc: '描述',
434+
parameterType: '参数类型',
435+
},
436+
},
423437
},
424438
compare: {
425439
is_null: '为空',

ui/src/locales/lang/zh-Hant/views/application-workflow.ts

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -393,6 +393,20 @@ export default {
393393
placeholder: '請輸入表達式',
394394
},
395395
},
396+
parameterExtractionNode: {
397+
label: 'Parameter Extraction',
398+
text: 'Extract structured parameters using AI model',
399+
result: 'Result',
400+
selectVariables: {
401+
label: 'Select Variables',
402+
placeholder: 'Please select variables',
403+
},
404+
extractParameters: {
405+
label: 'Extract Parameters',
406+
desc: 'Description',
407+
parameterType: 'Parameter Type',
408+
},
409+
},
396410
},
397411
compare: {
398412
is_null: '為空',

0 commit comments

Comments
 (0)