diff --git a/cognee/tasks/temporal_graph/enrich_events.py b/cognee/tasks/temporal_graph/enrich_events.py
index ef93da462..a43c8c4e9 100644
--- a/cognee/tasks/temporal_graph/enrich_events.py
+++ b/cognee/tasks/temporal_graph/enrich_events.py
@@ -1,6 +1,6 @@
 from typing import List
 
-from cognee.infrastructure.llm import LLMGateway
+from cognee.infrastructure.llm.extraction import extract_event_entities
 from cognee.modules.engine.models import Event
 from cognee.tasks.temporal_graph.models import EventWithEntities, EventEntityList
 
@@ -29,6 +29,6 @@ async def enrich_events(events: List[Event]) -> List[EventWithEntities]:
     events_json_str = json.dumps(events_json)
 
     # Extract entities from events
-    entity_result = await LLMGateway.extract_event_entities(events_json_str, EventEntityList)
+    entity_result = await extract_event_entities(events_json_str, EventEntityList)
 
     return entity_result.events
diff --git a/cognee/tasks/temporal_graph/extract_events_and_entities.py b/cognee/tasks/temporal_graph/extract_events_and_entities.py
index 8babc0ee5..f1c11bddf 100644
--- a/cognee/tasks/temporal_graph/extract_events_and_entities.py
+++ b/cognee/tasks/temporal_graph/extract_events_and_entities.py
@@ -1,6 +1,6 @@
 import asyncio
 from typing import Type, List
-from cognee.infrastructure.llm.LLMGateway import LLMGateway
+from cognee.infrastructure.llm.extraction import extract_event_graph
 from cognee.modules.chunking.models import DocumentChunk
 from cognee.tasks.temporal_graph.models import EventList
 from cognee.modules.engine.utils.generate_event_datapoint import generate_event_datapoint
@@ -21,7 +21,7 @@ async def extract_events_and_timestamps(data_chunks: List[DocumentChunk]) -> Lis
         List[DocumentChunk]: The same list of document chunks, enriched with extracted Event datapoints.
     """
     events = await asyncio.gather(
-        *[LLMGateway.extract_event_graph(chunk.text, EventList) for chunk in data_chunks]
+        *[extract_event_graph(chunk.text, EventList) for chunk in data_chunks]
     )
 
     for data_chunk, event_list in zip(data_chunks, events):