diff --git a/src/basic_memory/services/entity_service.py b/src/basic_memory/services/entity_service.py
index 1ab10528..39421474 100644
--- a/src/basic_memory/services/entity_service.py
+++ b/src/basic_memory/services/entity_service.py
@@ -11,7 +11,7 @@ from basic_memory.services.exceptions import EntityNotFoundError
 from .service import BaseService
 
 
-def entity_model(entity):
+def entity_model(entity: EntitySchema):
     model = EntityModel(
         name=entity.name,
         entity_type=entity.entity_type,