gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

refactor: Update LLMGateway usage

Browse source
This commit is contained in:
Igor Ilic 2025-09-09 14:05:40 +02:00
parent 85cda9f5c9
commit 9ab766fd86
2 changed files with 4 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
cognee/tasks/temporal_graph/enrich_events.py
View file

@ -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

4
cognee/tasks/temporal_graph/extract_events_and_entities.py
View file

@ -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):