diff --git a/custom_components/python_script/__init__.py b/custom_components/python_script/__init__.py
index fc8ce00..517ef15 100644
--- a/custom_components/python_script/__init__.py
+++ b/custom_components/python_script/__init__.py
@@ -86,7 +86,7 @@ def handler(call: ServiceCall) -> ServiceResponse:
         else:
             _LOGGER.debug("Load code from cache")
 
-        return execute_script(hass, call.data, _LOGGER, code)
+        return execute_script(hass, call.data, call.context, _LOGGER, code)
 
     hass.services.async_register(
         DOMAIN,
@@ -99,7 +99,7 @@ def handler(call: ServiceCall) -> ServiceResponse:
     return True
 
 
-def execute_script(hass: HomeAssistant, data: dict, logger, code) -> ServiceResponse:
+def execute_script(hass, data, context, logger, code) -> ServiceResponse:
     try:
         _LOGGER.debug("Run python script")
         vars = {**globals(), **locals()}