Update app.py

app.py

@@ -331,6 +331,9 @@ with gr.Blocks(title="CrystalGAT") as demo:
         # 初始化分子状态
         molecule_state = gr.State(init_molecule())
         
+        # 状态消息
+        status_msg = gr.Textbox(label="状态", interactive=False, value="请添加原子开始构建分子")
+        
         with gr.Row():
             # 原子选择
             with gr.Column():
@@ -398,27 +401,31 @@ with gr.Blocks(title="CrystalGAT") as demo:
         brittle_text = gr.Text(label="Brittle")
         brittle_img = gr.Image(type="pil", label="Brittle attention visualization")
     
-    # 更新原子列表显示
+    # 更新原子选择下拉菜单
+    def update_atom_dropdowns(molecule):
+        atom_choices = [f"{atom['id']}: {atom['type']}" for atom in molecule["atoms"]]
+        return gr.Dropdown.update(choices=atom_choices), gr.Dropdown.update(choices=atom_choices)
+    
+    # 更新原子列表
     def update_atoms_list(molecule):
-        atoms_data = [[i, atom["type"]] for i, atom in enumerate(molecule["atoms"])]
+        atoms_data = [[atom["id"], atom["type"]] for atom in molecule["atoms"]]
         return atoms_data
     
-    # 更新键列表显示
+    # 更新键列表
     def update_bonds_list(molecule):
         bonds_data = []
         for bond in molecule["bonds"]:
+            atom1_id = bond["atom1"]
+            atom2_id = bond["atom2"]
+            atom1_type = next((a["type"] for a in molecule["atoms"] if a["id"] == atom1_id), "?")
+            atom2_type = next((a["type"] for a in molecule["atoms"] if a["id"] == atom2_id), "?")
             bonds_data.append([
-                bond["atom1"],
-                bond["atom2"],
+                f"{atom1_id}: {atom1_type}",
+                f"{atom2_id}: {atom2_type}",
                 bond["type"]
             ])
         return bonds_data
     
-    # 更新原子选择下拉菜单
-    def update_atom_dropdowns(molecule):
-        atom_choices = [f"原子 {i}: {atom['type']}" for i, atom in enumerate(molecule["atoms"])]
-        return gr.Dropdown.update(choices=atom_choices), gr.Dropdown.update(choices=atom_choices)
-    
     # 事件处理
     add_atom_btn.click(
         fn=lambda atom, mol: add_atom(mol, atom),
@@ -432,17 +439,15 @@ with gr.Blocks(title="CrystalGAT") as demo:
         fn=update_atom_dropdowns,
         inputs=molecule_state,
         outputs=[atom1_select, atom2_select]
+    ).then(
+        fn=lambda: "原子添加成功!",
+        outputs=status_msg
     )
     
     add_bond_btn.click(
-        fn=lambda atom1_str, atom2_str, bond, mol: add_bond(
-            mol, 
-            int(atom1_str.split(":")[0].split(" ")[1]),  # 提取原子ID
-            int(atom2_str.split(":")[0].split(" ")[1]),  # 提取原子ID
-            bond
-        ),
+        fn=lambda atom1, atom2, bond, mol: (add_bond(mol, int(atom1.split(":")[0]), int(atom2.split(":")[0]), bond)) if atom1 and atom2 else (mol, "请先选择两个原子!"),
         inputs=[atom1_select, atom2_select, bond_select, molecule_state],
-        outputs=molecule_state
+        outputs=[molecule_state, status_msg]
     ).then(
         fn=update_bonds_list,
         inputs=molecule_state,
@@ -458,6 +463,9 @@ with gr.Blocks(title="CrystalGAT") as demo:
     ).then(
         fn=lambda: (gr.Dropdown.update(choices=[]), gr.Dropdown.update(choices=[])),
         outputs=[atom1_select, atom2_select]
+    ).then(
+        fn=lambda: "已清除所有原子和键",
+        outputs=status_msg
     )
     
     generate_btn.click(
@@ -468,6 +476,9 @@ with gr.Blocks(title="CrystalGAT") as demo:
         fn=visualize_molecule,
         inputs=molecule_state,
         outputs=molecule_img
+    ).then(
+        fn=lambda: "分子生成完成!",
+        outputs=status_msg
     )
     
     # 设置交互