diff --git a/internal/notification/handler.go b/internal/notification/handler.go
index 9ba7c00..76afc93 100644
--- a/internal/notification/handler.go
+++ b/internal/notification/handler.go
@@ -4,6 +4,7 @@ import (
 	"github.com/asaskevich/govalidator"
 	"github.com/labstack/echo/v4"
 
+	"tm/internal/customer"
 	"tm/internal/user"
 	"tm/pkg/logger"
 	"tm/pkg/response"
@@ -230,7 +231,7 @@ func (h *NotificationHandler) ViewNotification(c echo.Context) error {
 // @Security BearerAuth
 // @Router /api/v1/notifications [get]
 func (h *NotificationHandler) CustomerNotifications(c echo.Context) error {
-	userID, err := user.GetUserIDFromContext(c)
+	userID, err := customer.GetCustomerIDFromContext(c)
 	if err != nil {
 		return response.BadRequest(c, "User ID required", "User must be associated with a user")
 	}
@@ -274,7 +275,7 @@ func (h *NotificationHandler) CustomerNotifications(c echo.Context) error {
 func (h *NotificationHandler) PublicMarkSeen(c echo.Context) error {
 	notificationID := c.Param("id")
 
-	userID, err := user.GetUserIDFromContext(c)
+	userID, err := customer.GetCustomerIDFromContext(c)
 	if err != nil {
 		return response.BadRequest(c, "User ID required", "User must be associated with a user")
 	}