diff --git a/action.yml b/action.yml
index bf263a2..c77ed89 100644
--- a/action.yml
+++ b/action.yml
@@ -24,6 +24,11 @@ inputs:
     description: "Set to true to enable debug mode."
     default: 'false'
 
+  max_impacted_downstreams:
+    description: "Maximum number of downstream generations."
+    required: false
+    default: 30
+
 runs:
   using: "composite"
   steps:
@@ -46,6 +51,7 @@ runs:
         GITHUB_ACTION_PATH: ${{ github.action_path }}
         GITHUB_BASE_REF: ${{ github.base_ref }}
         DEBUG_MODE: ${{ inputs.debug_mode }}
+        MAX_IMPACTED_DOWNSTREAMS: ${{ inputs.max_impacted_downstreams }}
 
     # Post a comment on the PR.
     - uses: marocchino/sticky-pull-request-comment@v2.5.0
diff --git a/src/impact_analysis.py b/src/impact_analysis.py
index 38a24ad..23d43b0 100644
--- a/src/impact_analysis.py
+++ b/src/impact_analysis.py
@@ -9,12 +9,12 @@
 from datahub.ingestion.graph.client import DatahubClientConfig, DataHubGraph
 from datahub.metadata.schema_classes import DatasetPropertiesClass
 from datahub.telemetry import telemetry
-from os import environ
+
 from rendering import datahub_url_from_urn, format_entity
 
 OUTPUT_PATH = pathlib.Path("impact_analysis.md")
 DBT_ID_PROP = "dbt_unique_id"
-MAX_IMPACTED_DOWNSTREAMS = environ.get("MAX_IMPACTED_DOWNSTREAMS") or 30
+MAX_IMPACTED_DOWNSTREAMS = os.getenv("MAX_IMPACTED_DOWNSTREAMS")
 MAX_DOWNSTREAMS_TO_FETCH = 1000