From d764537143aa44049b2e3f5539d53829ae4fcc85 Mon Sep 17 00:00:00 2001
From: carry <luokairui@carry.fit>
Date: Wed, 16 Apr 2025 12:39:48 +0800
Subject: [PATCH] =?UTF-8?q?feat(dataset=5Fgenerate=5Fpage):=20=E6=9B=B4?=
 =?UTF-8?q?=E6=96=B0=E6=95=B0=E6=8D=AE=E9=9B=86=E7=94=9F=E6=88=90=E9=A1=B5?=
 =?UTF-8?q?=E9=9D=A2=E5=8A=9F=E8=83=BD?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- 添加模板变量列表展示和编辑功能
- 实现模板选择后动态更新变量列表
- 增加生成数据集按钮和相关逻辑
- 优化页面布局和交互
---
 frontend/dataset_generate_page.py | 69 +++++++++++++++++++++++--------
 1 file changed, 51 insertions(+), 18 deletions(-)

diff --git a/frontend/dataset_generate_page.py b/frontend/dataset_generate_page.py
index a264b3c..dd57cb6 100644
--- a/frontend/dataset_generate_page.py
+++ b/frontend/dataset_generate_page.py
@@ -11,51 +11,84 @@ def dataset_generate_page():
         gr.Markdown("## 数据集生成")
         with gr.Row():
             with gr.Column():
-                # 获取文档列表并设置初始值
                 docs_list = [str(doc.name) for doc in get_docs()]
                 initial_doc = docs_list[0] if docs_list else None
-                
                 doc_dropdown = gr.Dropdown(
                     choices=docs_list,
-                    value=initial_doc,  # 设置初始选中项
+                    value=initial_doc,
                     label="选择文档",
-                    allow_custom_value=True,
                     interactive=True
                 )
-                doc_state = gr.State(value=initial_doc)  # 用文档初始值初始化状态
+                doc_state = gr.State(value=initial_doc)
 
             with gr.Column():
-                # 获取模板列表并设置初始值
                 prompts = get_prompt_store().all()
                 prompt_list = [f"{p['id']} {p['name']}" for p in prompts]
                 initial_prompt = prompt_list[0] if prompt_list else None
                 
+                # 初始化Dataframe的值
+                initial_dataframe_value = []
+                if initial_prompt:
+                    selected_prompt_id = int(initial_prompt.split(" ")[0])
+                    prompt_data = get_prompt_store().get(doc_id=selected_prompt_id)
+                    prompt_content = prompt_data["content"]
+                    prompt_template = PromptTemplate.from_template(prompt_content)
+                    input_variables = prompt_template.input_variables
+                    initial_dataframe_value = [[var, ""] for var in input_variables]
+
                 prompt_dropdown = gr.Dropdown(
                     choices=prompt_list,
-                    value=initial_prompt,  # 设置初始选中项
+                    value=initial_prompt,
                     label="选择模板",
-                    allow_custom_value=True,
                     interactive=True
                 )
-                prompt_state = gr.State(value=initial_prompt)  # 用模板初始值初始化状态
+                prompt_state = gr.State(value=initial_prompt)
+
+        generate_button = gr.Button("生成数据集")
+
+        variables_dataframe = gr.Dataframe(
+            headers=["变量名", "变量值"],
+            datatype=["str", "str"],
+            interactive=True,
+            label="变量列表",
+            value=initial_dataframe_value  # 设置初始化数据
+        )
+
+        output_text = gr.Textbox(label="生成结果", interactive=False)
 
-        # 新增事件绑定
         def on_doc_change(selected_doc):
-            print(f"文档选择已更改为: {selected_doc}")
+            # print(f"文档选择已更改为: {selected_doc}")
             return selected_doc
-
         def on_prompt_change(selected_prompt):
-            print(f"模板选择已更改为: {selected_prompt}")
+            if not selected_prompt:
+                return None, []
             selected_prompt_id = int(selected_prompt.split(" ")[0])
-            selected_prompt = get_prompt_store().get(doc_id=selected_prompt_id)
-            selected_prompt = selected_prompt["content"]
-            selected_prompt = PromptTemplate.from_template(selected_prompt)
-            input_variables_list = selected_prompt.input_variables
-            return selected_prompt
+            prompt_data = get_prompt_store().get(doc_id=selected_prompt_id)
+            prompt_content = prompt_data["content"]
+            prompt_template = PromptTemplate.from_template(prompt_content)
+            input_variables = prompt_template.input_variables
+            dataframe_value = [[var, ""] for var in input_variables]
+            return selected_prompt, dataframe_value
 
+        def on_generate_click(doc_state, prompt_state, variables_dataframe):
+            variables_dict = {}
+            # 正确遍历DataFrame的行数据
+            for _, row in variables_dataframe.iterrows():
+                var_name = row['变量名'].strip()
+                var_value = row['变量值'].strip()
+                if var_name:
+                    variables_dict[var_name] = var_value
+
+        
         doc_dropdown.change(on_doc_change, inputs=doc_dropdown, outputs=doc_state)
         prompt_dropdown.change(on_prompt_change, inputs=prompt_dropdown, outputs=prompt_state)
 
+        generate_button.click(
+            on_generate_click,
+            inputs=[doc_state, prompt_state, variables_dataframe],
+            outputs=output_text
+        )
+
     return demo
 
 if __name__ == "__main__":