diff --git a/scrapegraphai/graphs/code_generator_graph.py b/scrapegraphai/graphs/code_generator_graph.py
index c4fc81e7..c4fd6d7a 100644
--- a/scrapegraphai/graphs/code_generator_graph.py
+++ b/scrapegraphai/graphs/code_generator_graph.py
@@ -12,6 +12,7 @@ from ..nodes import (
     GenerateAnswerNode,
     PromptRefinerNode,
     HtmlAnalyzerNode,
+    GenerateCodeNode,
 )
 
 class CodeGeneratorGraph(AbstractGraph):
@@ -115,7 +116,7 @@ class CodeGeneratorGraph(AbstractGraph):
         
         generate_code_node = GenerateCodeNode(
             input="user_prompt & refined_prompt & html_info & doc & answer",
-            output=["code"],
+            output=["generated_code"],
             node_config={
                 "llm_model": self.llm_model,
                 "additional_info": self.config.get("additional_info"),
diff --git a/scrapegraphai/nodes/__init__.py b/scrapegraphai/nodes/__init__.py
index a5ffb757..e5427044 100644
--- a/scrapegraphai/nodes/__init__.py
+++ b/scrapegraphai/nodes/__init__.py
@@ -24,4 +24,5 @@ from .fetch_screen_node import FetchScreenNode
 from .generate_answer_from_image_node import GenerateAnswerFromImageNode
 from .concat_answers_node import ConcatAnswersNode
 from .prompt_refiner_node import PromptRefinerNode
-from .html_analyzer_node import HtmlAnalyzerNode
\ No newline at end of file
+from .html_analyzer_node import HtmlAnalyzerNode
+from .generate_code_node import GenerateCodeNode
\ No newline at end of file