ruff format

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]:

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
+from .extract_event_entities import extract_event_entities

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(
+async def extract_event_entities(content: str, response_model: Type[BaseModel]):
-    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
+    return content_graph

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
+    return content_graph

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"]}
+    metadata: dict = {"index_fields": ["name"]}

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(...)
+    time_to: Timestamp = Field(...)

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(...)
+    timestamp_str: str = Field(...)

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)
+    return Event(**event_data)

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
+    return time

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
-

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]
+    return entity_types[entity_type_name]

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
+    return entity_result.events

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
+    return data_chunks

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
+    return data_chunks

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]
+    events: List[EventWithEntities]

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())
+        loop.run_until_complete(loop.shutdown_asyncgens())