From 58a3be7c126b2d3d14ae47e53d2891ae4d12cd5b Mon Sep 17 00:00:00 2001 From: hajdul88 <52442977+hajdul88@users.noreply.github.com> Date: Wed, 27 Aug 2025 18:04:58 +0200 Subject: [PATCH] ruff format --- cognee/api/v1/cognify/cognify.py | 6 +++++- .../litellm_instructor/extraction/__init__.py | 2 +- .../extraction/extract_event_entities.py | 6 ++---- .../extraction/knowledge_graph/extract_event_graph.py | 3 ++- cognee/modules/engine/models/Event.py | 2 +- cognee/modules/engine/models/Interval.py | 3 ++- cognee/modules/engine/models/Timestamp.py | 2 +- cognee/modules/engine/utils/generate_event_datapoint.py | 3 ++- .../modules/engine/utils/generate_timestamp_datapoint.py | 4 +++- cognee/tasks/temporal_graph/__init__.py | 1 - cognee/tasks/temporal_graph/add_entities_to_event.py | 4 +++- cognee/tasks/temporal_graph/enrich_events.py | 5 +++-- cognee/tasks/temporal_graph/extract_events_and_entities.py | 2 +- .../temporal_graph/extract_knowledge_graph_from_events.py | 7 +++++-- cognee/tasks/temporal_graph/models.py | 3 +-- examples/python/temporal_example.py | 2 +- 16 files changed, 33 insertions(+), 22 deletions(-) diff --git a/cognee/api/v1/cognify/cognify.py b/cognee/api/v1/cognify/cognify.py index dee4e79be..a0803ff96 100644 --- a/cognee/api/v1/cognify/cognify.py +++ b/cognee/api/v1/cognify/cognify.py @@ -22,7 +22,10 @@ from cognee.tasks.graph import extract_graph_from_data from cognee.tasks.storage import add_data_points from cognee.tasks.summarization import summarize_text from cognee.modules.pipelines.layers.pipeline_execution_mode import get_pipeline_executor -from cognee.tasks.temporal_graph import extract_events_and_timestamps, extract_knowledge_graph_from_events +from cognee.tasks.temporal_graph import ( + extract_events_and_timestamps, + extract_knowledge_graph_from_events, +) logger = get_logger("cognify") @@ -230,6 +233,7 @@ async def get_default_tasks( # TODO: Find out a better way to do this (Boris's return default_tasks + async def get_temporal_tasks( user: User = None, chunker=TextChunker, chunk_size: int = None ) -> list[Task]: diff --git a/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/extraction/__init__.py b/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/extraction/__init__.py index 24006c046..72e3c755f 100644 --- a/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/extraction/__init__.py +++ b/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/extraction/__init__.py @@ -2,4 +2,4 @@ from .knowledge_graph.extract_content_graph import extract_content_graph from .knowledge_graph.extract_event_graph import extract_event_graph from .extract_categories import extract_categories from .extract_summary import extract_summary, extract_code_summary -from .extract_event_entities import extract_event_entities \ No newline at end of file +from .extract_event_entities import extract_event_entities diff --git a/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/extraction/extract_event_entities.py b/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/extraction/extract_event_entities.py index 123c05269..ad33863b0 100644 --- a/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/extraction/extract_event_entities.py +++ b/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/extraction/extract_event_entities.py @@ -7,9 +7,7 @@ from cognee.infrastructure.llm.config import ( ) -async def extract_event_entities( - content: str, response_model: Type[BaseModel] -): +async def extract_event_entities(content: str, response_model: Type[BaseModel]): """Extract event entities from content using LLM.""" llm_config = get_llm_config() @@ -30,4 +28,4 @@ async def extract_event_entities( content, system_prompt, response_model ) - return content_graph \ No newline at end of file + return content_graph diff --git a/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/extraction/knowledge_graph/extract_event_graph.py b/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/extraction/knowledge_graph/extract_event_graph.py index 2a0c0cab8..0373649f2 100644 --- a/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/extraction/knowledge_graph/extract_event_graph.py +++ b/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/extraction/knowledge_graph/extract_event_graph.py @@ -7,6 +7,7 @@ from cognee.infrastructure.llm.config import ( get_llm_config, ) + async def extract_event_graph( content: str, response_model: Type[BaseModel], system_prompt: str = None ): @@ -31,4 +32,4 @@ async def extract_event_graph( content, system_prompt, response_model ) - return content_graph \ No newline at end of file + return content_graph diff --git a/cognee/modules/engine/models/Event.py b/cognee/modules/engine/models/Event.py index 88141e602..4a0bab830 100644 --- a/cognee/modules/engine/models/Event.py +++ b/cognee/modules/engine/models/Event.py @@ -13,4 +13,4 @@ class Event(DataPoint): location: Optional[str] = None attributes: SkipValidation[Any] = None - metadata: dict = {"index_fields": ["name"]} \ No newline at end of file + metadata: dict = {"index_fields": ["name"]} diff --git a/cognee/modules/engine/models/Interval.py b/cognee/modules/engine/models/Interval.py index 3666bf69d..914bc62ea 100644 --- a/cognee/modules/engine/models/Interval.py +++ b/cognee/modules/engine/models/Interval.py @@ -2,6 +2,7 @@ from pydantic import Field from cognee.infrastructure.engine import DataPoint from cognee.modules.engine.models.Timestamp import Timestamp + class Interval(DataPoint): time_from: Timestamp = Field(...) - time_to: Timestamp = Field(...) \ No newline at end of file + time_to: Timestamp = Field(...) diff --git a/cognee/modules/engine/models/Timestamp.py b/cognee/modules/engine/models/Timestamp.py index 38977c348..31779683a 100644 --- a/cognee/modules/engine/models/Timestamp.py +++ b/cognee/modules/engine/models/Timestamp.py @@ -10,4 +10,4 @@ class Timestamp(DataPoint): hour: int = Field(...) minute: int = Field(...) second: int = Field(...) - timestamp_str: str = Field(...) \ No newline at end of file + timestamp_str: str = Field(...) diff --git a/cognee/modules/engine/utils/generate_event_datapoint.py b/cognee/modules/engine/utils/generate_event_datapoint.py index aeec325d9..cc56763ae 100644 --- a/cognee/modules/engine/utils/generate_event_datapoint.py +++ b/cognee/modules/engine/utils/generate_event_datapoint.py @@ -1,6 +1,7 @@ from cognee.modules.engine.models import Interval, Event from cognee.modules.engine.utils.generate_timestamp_datapoint import generate_timestamp_datapoint + def generate_event_datapoint(event) -> Event: """Create an Event datapoint from an event model.""" # Base event data @@ -27,4 +28,4 @@ def generate_event_datapoint(event) -> Event: temporal_info = f"\n---\nTime data: {timestamp.timestamp_str}" event_data["description"] = (event_data["description"] or "Event") + temporal_info - return Event(**event_data) \ No newline at end of file + return Event(**event_data) diff --git a/cognee/modules/engine/utils/generate_timestamp_datapoint.py b/cognee/modules/engine/utils/generate_timestamp_datapoint.py index cbef2d177..6f2cdf6d1 100644 --- a/cognee/modules/engine/utils/generate_timestamp_datapoint.py +++ b/cognee/modules/engine/utils/generate_timestamp_datapoint.py @@ -2,6 +2,7 @@ from datetime import datetime, timezone from cognee.modules.engine.models import Interval, Timestamp, Event from cognee.modules.engine.utils import generate_node_id + def generate_timestamp_datapoint(ts: Timestamp) -> Timestamp: """Create a Timestamp datapoint from a Timestamp model.""" time_at = date_to_int(ts) @@ -20,8 +21,9 @@ def generate_timestamp_datapoint(ts: Timestamp) -> Timestamp: timestamp_str=timestamp_str, ) + def date_to_int(ts: Timestamp) -> int: """Convert timestamp to integer milliseconds.""" dt = datetime(ts.year, ts.month, ts.day, ts.hour, ts.minute, ts.second, tzinfo=timezone.utc) time = int(dt.timestamp() * 1000) - return time \ No newline at end of file + return time diff --git a/cognee/tasks/temporal_graph/__init__.py b/cognee/tasks/temporal_graph/__init__.py index 991553605..11d812541 100644 --- a/cognee/tasks/temporal_graph/__init__.py +++ b/cognee/tasks/temporal_graph/__init__.py @@ -1,3 +1,2 @@ from .extract_events_and_entities import extract_events_and_timestamps from .extract_knowledge_graph_from_events import extract_knowledge_graph_from_events - diff --git a/cognee/tasks/temporal_graph/add_entities_to_event.py b/cognee/tasks/temporal_graph/add_entities_to_event.py index 5585a1b50..2cb4b1425 100644 --- a/cognee/tasks/temporal_graph/add_entities_to_event.py +++ b/cognee/tasks/temporal_graph/add_entities_to_event.py @@ -5,6 +5,7 @@ from cognee.modules.engine.models.EntityType import EntityType from cognee.infrastructure.engine.models.Edge import Edge from cognee.modules.engine.utils import generate_node_id, generate_node_name + def add_entities_to_event(event: Event, event_with_entities: EventWithEntities) -> None: """Add entities to event via attributes field.""" if not event_with_entities.attributes: @@ -38,6 +39,7 @@ def add_entities_to_event(event: Event, event_with_entities: EventWithEntities) event.attributes = [] event.attributes.append((edge, [entity])) + def get_or_create_entity_type(entity_types: dict, entity_type_name: str) -> EntityType: """Get existing entity type or create new one.""" if entity_type_name not in entity_types: @@ -52,4 +54,4 @@ def get_or_create_entity_type(entity_types: dict, entity_type_name: str) -> Enti ) entity_types[entity_type_name] = entity_type - return entity_types[entity_type_name] \ No newline at end of file + return entity_types[entity_type_name] diff --git a/cognee/tasks/temporal_graph/enrich_events.py b/cognee/tasks/temporal_graph/enrich_events.py index 4c9edb2bb..bedd642eb 100644 --- a/cognee/tasks/temporal_graph/enrich_events.py +++ b/cognee/tasks/temporal_graph/enrich_events.py @@ -2,7 +2,8 @@ from typing import List from cognee.infrastructure.llm import LLMGateway from cognee.modules.engine.models import Event -from cognee.tasks.temporal_graph.models import EventWithEntities,EventEntityList +from cognee.tasks.temporal_graph.models import EventWithEntities, EventEntityList + async def enrich_events(events: List[Event]) -> List[EventWithEntities]: """Extract entities from events and return enriched events.""" @@ -18,4 +19,4 @@ async def enrich_events(events: List[Event]) -> List[EventWithEntities]: # Extract entities from events entity_result = await LLMGateway.extract_event_entities(events_json_str, EventEntityList) - return entity_result.events \ No newline at end of file + 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 bf4367f6a..de0cdd601 100644 --- a/cognee/tasks/temporal_graph/extract_events_and_entities.py +++ b/cognee/tasks/temporal_graph/extract_events_and_entities.py @@ -17,4 +17,4 @@ async def extract_events_and_timestamps(data_chunks: List[DocumentChunk]) -> Lis event_datapoint = generate_event_datapoint(event) data_chunk.contains.append(event_datapoint) - return data_chunks \ No newline at end of file + return data_chunks diff --git a/cognee/tasks/temporal_graph/extract_knowledge_graph_from_events.py b/cognee/tasks/temporal_graph/extract_knowledge_graph_from_events.py index 0e49c5296..8cbcc3c22 100644 --- a/cognee/tasks/temporal_graph/extract_knowledge_graph_from_events.py +++ b/cognee/tasks/temporal_graph/extract_knowledge_graph_from_events.py @@ -4,7 +4,10 @@ from cognee.modules.engine.models import Event from cognee.tasks.temporal_graph.enrich_events import enrich_events from cognee.tasks.temporal_graph.add_entities_to_event import add_entities_to_event -async def extract_knowledge_graph_from_events(data_chunks: List[DocumentChunk]) -> List[DocumentChunk]: + +async def extract_knowledge_graph_from_events( + data_chunks: List[DocumentChunk], +) -> List[DocumentChunk]: """Extract events from chunks and enrich them with entities.""" # Extract events from chunks all_events = [] @@ -23,4 +26,4 @@ async def extract_knowledge_graph_from_events(data_chunks: List[DocumentChunk]) for event, enriched_event in zip(all_events, enriched_events): add_entities_to_event(event, enriched_event) - return data_chunks \ No newline at end of file + return data_chunks diff --git a/cognee/tasks/temporal_graph/models.py b/cognee/tasks/temporal_graph/models.py index 3818110c5..ef5cd42c9 100644 --- a/cognee/tasks/temporal_graph/models.py +++ b/cognee/tasks/temporal_graph/models.py @@ -2,7 +2,6 @@ from typing import Optional, List from pydantic import BaseModel, Field - class Timestamp(BaseModel): year: int = Field(..., ge=1, le=9999) month: int = Field(..., ge=1, le=12) @@ -47,4 +46,4 @@ class EventWithEntities(BaseModel): class EventEntityList(BaseModel): - events: List[EventWithEntities] \ No newline at end of file + events: List[EventWithEntities] diff --git a/examples/python/temporal_example.py b/examples/python/temporal_example.py index c61c80ac4..c4c1c9875 100644 --- a/examples/python/temporal_example.py +++ b/examples/python/temporal_example.py @@ -116,4 +116,4 @@ if __name__ == "__main__": try: loop.run_until_complete(main()) finally: - loop.run_until_complete(loop.shutdown_asyncgens()) \ No newline at end of file + loop.run_until_complete(loop.shutdown_asyncgens())