diff --git a/servers/structure_generate/server.py b/servers/structure_generate/server.py
index 6519e9c..4a48f02 100644
--- a/servers/structure_generate/server.py
+++ b/servers/structure_generate/server.py
@@ -24,6 +24,7 @@
 from pymatgen.core.structure import Structure
 from pymatgen.core.surface import SlabGenerator
 from pymatgen.analysis.local_env import CrystalNN
+from pymatgen.core import Lattice
 
 # Local/custom imports
 from dp.agent.server import CalculationMCPServer