feat: extract entities from events

cognee/modules/chunking/models/DocumentChunk.py

@@ -3,6 +3,8 @@ from typing import List
 from cognee.infrastructure.engine import DataPoint
 from cognee.modules.data.processing.document_types import Document
 from cognee.modules.engine.models import Entity
+from typing import Union
+from cognee.temporal_poc.models.models import Event
 
 
 class DocumentChunk(DataPoint):
@@ -30,6 +32,6 @@ class DocumentChunk(DataPoint):
     chunk_index: int
     cut_type: str
     is_part_of: Document
-    contains: List[Entity] = None
+    contains: List[Union[Entity, Event]] = None
 
     metadata: dict = {"index_fields": ["text"]}

cognee/tasks/storage/index_data_points.py

@@ -1,9 +1,11 @@
+import asyncio
 from cognee.shared.logging_utils import get_logger
 
 from cognee.infrastructure.databases.exceptions.EmbeddingException import EmbeddingException
 from cognee.infrastructure.databases.vector import get_vector_engine
 from cognee.infrastructure.engine import DataPoint
 
+vector_index_lock = asyncio.Lock()
 logger = get_logger("index_data_points")
 
 
@@ -22,9 +24,10 @@ async def index_data_points(data_points: list[DataPoint]):
 
             index_name = f"{data_point_type.__name__}_{field_name}"
 
-            if index_name not in created_indexes:
-                await vector_engine.create_vector_index(data_point_type.__name__, field_name)
-                created_indexes[index_name] = True
+            async with vector_index_lock:
+                if index_name not in created_indexes:
+                    await vector_engine.create_vector_index(data_point_type.__name__, field_name)
+                    created_indexes[index_name] = True
 
             if index_name not in index_points:
                 index_points[index_name] = []

cognee/temporal_poc/datapoints/datapoints.py

@@ -0,0 +1,23 @@
+from cognee.infrastructure.engine import DataPoint
+from cognee.modules.engine.models.EntityType import EntityType
+from typing import Optional
+from pydantic import BaseModel, Field, ConfigDict
+
+
+class Timestamp(DataPoint):
+    time_at: int = Field(...)
+
+
+class Interval(DataPoint):
+    time_from: Timestamp = Field(...)
+    time_to: Timestamp = Field(...)
+
+
+class Event(DataPoint):
+    name: str
+    description: Optional[str] = None
+    at: Optional[Timestamp] = None
+    during: Optional[Interval] = None
+    location: Optional[str] = None
+
+    metadata: dict = {"index_fields": ["name"]}

cognee/temporal_poc/event_extraction.py

@@ -0,0 +1,126 @@
+import asyncio
+import uuid
+
+from pydantic import BaseModel
+from typing import Type, List
+from uuid import uuid5
+from datetime import datetime, timezone
+
+from cognee.infrastructure.llm.get_llm_client import get_llm_client
+from cognee.modules.chunking.models.DocumentChunk import DocumentChunk
+from cognee.temporal_poc.models.models import EventList
+from cognee.temporal_poc.datapoints.datapoints import Interval, Timestamp, Event
+
+
+# Global system prompt for event extraction
+EVENT_EXTRACTION_SYSTEM_PROMPT = """
+        For the purposes of building event-based knowledge graphs, you are tasked with extracting highly granular stream events from a text. The events are defined as follows:
+        ## Event Definition
+        - Anything with a date or a timestamp is an event
+        - Anything that took place in time (even if the time is unknown) is an event
+        - Anything that lasted over a period of time, or happened in an instant is an event: from historical milestones (wars, presidencies, olympiads) to personal milestones (birth, death, employment, etc.), to mundane actions (a walk, a conversation, etc.)
+        - **ANY action or verb represents an event** - this is the most important rule
+        - Every single verb in the text corresponds to an event that must be extracted
+        - This includes: thinking, feeling, seeing, hearing, moving, speaking, writing, reading, eating, sleeping, working, playing, studying, traveling, meeting, calling, texting, buying, selling, creating, destroying, building, breaking, starting, stopping, beginning, ending, etc.
+        - Even the most mundane or obvious actions are events: "he walked", "she sat", "they talked", "I thought", "we waited"
+        ## Requirements
+        - **Be extremely thorough** - extract EVERY event mentioned, no matter how small or obvious
+        - **Timestamped first" - every time stamp, or date should have atleast one event
+        - **Verbs/actions  = one event** - After you are done with timestamped events -- every verb that is an action should have a corresponding event.
+        - We expect long streams of events from any piece of text, easily reaching a hundred events
+        - Granularity and richness of the stream is key to our success and is of utmost importance
+        - Not all events will have timestamps, add timestamps only to known events
+        - For events that were instantaneous, just attach the time_from or time_to property don't create both
+        - **Do not skip any events** - if you're unsure whether something is an event, extract it anyway
+        - **Quantity over filtering** - it's better to extract too many events than to miss any
+        - **Descriptions** - Always include the event description together with entities (Who did what, what happened? What is the event?). If you can include the corresponding part from the text.
+        ## Output Format
+        Your reply should be a JSON: list of dictionaries with the following structure:
+        ```python
+        class Event(BaseModel):
+            name: str [concise]
+            description: Optional[str] = None
+            time_from: Optional[Timestamp] = None
+            time_to: Optional[Timestamp] = None
+            location: Optional[str] = None
+        ```
+"""
+
+
+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
+
+
+async def extract_event_graph(
+    content: str, response_model: Type[BaseModel], system_prompt: str = None
+):
+    """Extract event graph from content using LLM."""
+    llm_client = get_llm_client()
+
+    if system_prompt is None:
+        system_prompt = EVENT_EXTRACTION_SYSTEM_PROMPT
+
+    content_graph = await llm_client.acreate_structured_output(
+        content, system_prompt, response_model
+    )
+
+    return content_graph
+
+
+async def extract_events_and_entities(data_chunks: List[DocumentChunk]) -> List[DocumentChunk]:
+    """Extracts events and entities from a chunk of documents."""
+    events = await asyncio.gather(
+        *[extract_event_graph(chunk.text, EventList) for chunk in data_chunks]
+    )
+
+    for data_chunk, event_list in zip(data_chunks, events):
+        for event in event_list.events:
+            if event.time_from and event.time_to:
+                event_time_from = date_to_int(event.time_from)
+                event_time_to = date_to_int(event.time_to)
+                timestamp_time_from = Timestamp(
+                    id=uuid5(uuid.NAMESPACE_OID, name=str(event_time_from)), time_at=event_time_from
+                )
+                timestamp_time_to = Timestamp(
+                    id=uuid5(uuid.NAMESPACE_OID, name=str(event_time_to)), time_at=event_time_to
+                )
+                event_interval = Interval(time_from=timestamp_time_from, time_to=timestamp_time_to)
+                event_datapoint = Event(
+                    name=event.name,
+                    description=event.description,
+                    during=event_interval,
+                    location=event.location,
+                )
+            elif event.time_from:
+                event_time_from = date_to_int(event.time_from)
+                event_time_at = Timestamp(
+                    id=uuid5(uuid.NAMESPACE_OID, name=str(event_time_from)), time_at=event_time_from
+                )
+                event_datapoint = Event(
+                    name=event.name,
+                    description=event.description,
+                    at=event_time_at,
+                    location=event.location,
+                )
+            elif event.time_to:
+                event_time_to = date_to_int(event.time_to)
+                event_time_at = Timestamp(
+                    id=uuid5(uuid.NAMESPACE_OID, name=str(event_time_to)), time_at=event_time_to
+                )
+                event_datapoint = Event(
+                    name=event.name,
+                    description=event.description,
+                    at=event_time_at,
+                    location=event.location,
+                )
+            else:
+                event_datapoint = Event(
+                    name=event.name, description=event.description, location=event.location
+                )
+
+            data_chunk.contains.append(event_datapoint)
+
+    return data_chunks

cognee/temporal_poc/event_knowledge_graph.py

@@ -0,0 +1,119 @@
+from typing import List, Type
+from pydantic import BaseModel
+
+from cognee.infrastructure.llm.get_llm_client import get_llm_client
+from cognee.modules.chunking.models.DocumentChunk import DocumentChunk
+from cognee.modules.engine.models.Entity import Entity
+from cognee.modules.engine.models.EntityType import EntityType
+from cognee.temporal_poc.models.models import EventEntityList
+
+ENTITY_EXTRACTION_SYSTEM_PROMPT = """For the purposes of building event-based knowledge graphs, you are tasked with extracting highly granular entities from events text. An entity is any distinct, identifiable thing, person, place, object, organization, concept, or phenomenon that can be named, referenced, or described in the event context. This includes but is not limited to: people, places, objects, organizations, concepts, events, processes, states, conditions, properties, attributes, roles, functions, and any other meaningful referents that contribute to understanding the event.
+**Temporal Entity Exclusion**: Do not extract timestamp-like entities (dates, times, durations) as these are handled separately. However, extract named temporal periods, eras, historical epochs, and culturally significant time references
+## Input Format
+The input will be a list of dictionaries, each containing:
+- `event_name`: The name of the event
+- `description`: The description of the event
+
+## Task
+For each event, extract all entities mentioned in the event description and determine their relationship to the event.
+
+## Output Format
+Return the same enriched JSON with an additional key in each dictionary: `attributes`.
+
+The `attributes` should be a list of dictionaries, each containing:
+- `entity`: The name of the entity
+- `entity_type`: The type/category of the entity (person, place, organization, object, concept, etc.)
+- `relationship`: A concise description of how the entity relates to the event
+
+## Requirements
+- **Be extremely thorough** - extract EVERY non-temporal entity mentioned, no matter how small, obvious, or seemingly insignificant
+- **After you are done with obvious entities, every noun, pronoun, proper noun, and named reference =  one entity**
+- We expect rich entity networks from any event, easily reaching a dozens of entities per event
+- Granularity and richness of the entity extraction is key to our success and is of utmost importance
+- **Do not skip any entities** - if you're unsure whether something is an entity, extract it anyway
+- Use the event name for context when determining relationships
+- Relationships should be technical with one or at most two words. If two words, use underscore camelcase style
+- Relationships could imply general meaning like: subject, object, participant, recipient, agent, instrument, tool, source, cause, effect, purpose, manner, resource, etc.
+- You can combine two words to form a relationship name: subject_role, previous_owner, etc.
+- Focus on how the entity specifically relates to the event
+"""
+
+
+async def extract_event_entities(
+    content: str, response_model: Type[BaseModel], system_prompt: str = None
+):
+    """Extract event entities from content using LLM."""
+    llm_client = get_llm_client()
+
+    if system_prompt is None:
+        system_prompt = ENTITY_EXTRACTION_SYSTEM_PROMPT
+
+    content_graph = await llm_client.acreate_structured_output(
+        content, system_prompt, response_model
+    )
+
+    return content_graph
+
+
+async def extract_event_knowledge_graph(data_chunks: List[DocumentChunk]) -> List[DocumentChunk]:
+    """Extract events from chunks and process them for entity extraction."""
+    from cognee.temporal_poc.datapoints.datapoints import Event
+    import asyncio
+    import json
+
+    # Extract events from chunks - create a list of lists
+    chunk_events = [
+        [
+            {"event_name": item.name, "description": item.description or ""}
+            for item in chunk.contains
+            if isinstance(item, Event)
+        ]
+        for chunk in data_chunks
+    ]
+
+    # Convert each chunk's events to JSON and gather all calls
+    events_jsons = [json.dumps(events) if events else "[]" for events in chunk_events]
+
+    # Extract entities from all chunks concurrently
+    entity_results = await asyncio.gather(
+        *[extract_event_entities(events_json, EventEntityList) for events_json in events_jsons]
+    )
+
+    # Process entity results and enrich chunks
+    for chunk, entity_result in zip(data_chunks, entity_results):
+        add_entities_to_chunk(chunk, entity_result)
+
+    return data_chunks
+
+
+def add_entities_to_chunk(chunk: DocumentChunk, entity_result) -> None:
+    """Add entities and entity types to a chunk."""
+    entity_types = {}  # Cache to avoid duplicates
+    for event_with_entities in entity_result.events:
+        for attribute in event_with_entities.attributes:
+            entity_type = get_or_create_entity_type(entity_types, attribute.entity_type)
+            create_entity(chunk, attribute, entity_type)
+
+
+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:
+        entity_type = EntityType(name=entity_type_name, description=f"Type for {entity_type_name}")
+        entity_types[entity_type_name] = entity_type
+
+    return entity_types[entity_type_name]
+
+
+def create_entity(chunk: DocumentChunk, attribute, entity_type: EntityType) -> None:
+    """Create and add entity to chunk."""
+    entity = Entity(
+        name=attribute.entity,
+        is_a=entity_type,
+        description=f"Entity {attribute.entity} of type {attribute.entity_type}",
+    )
+    chunk.contains.append(entity)
+
+
+async def process_event_knowledge_graph(data_chunks: List[DocumentChunk]) -> List[DocumentChunk]:
+    """Process document chunks for event knowledge graph construction."""
+    return await extract_event_knowledge_graph(data_chunks)

cognee/temporal_poc/models/models.py

@@ -0,0 +1,45 @@
+from typing import Optional, Literal, List
+from pydantic import BaseModel, Field, root_validator, ValidationError
+from cognee.modules.engine.models.Entity import Entity
+
+
+class Timestamp(BaseModel):
+    year: int = Field(..., ge=1, le=9999)
+    month: int = Field(..., ge=1, le=12)
+    day: int = Field(..., ge=1, le=31)
+    hour: int = Field(..., ge=0, le=23)
+    minute: int = Field(..., ge=0, le=59)
+    second: int = Field(..., ge=0, le=59)
+
+
+class Interval(BaseModel):
+    starts_at: Timestamp
+    ends_at: Timestamp
+
+
+class Event(BaseModel):
+    name: str
+    description: Optional[str] = None
+    time_from: Optional[Timestamp] = None
+    time_to: Optional[Timestamp] = None
+    location: Optional[str] = None
+
+
+class EventList(BaseModel):
+    events: List[Event]
+
+
+class EntityAttribute(BaseModel):
+    entity: str
+    entity_type: str
+    relationship: str
+
+
+class EventWithEntities(BaseModel):
+    event_name: str
+    description: Optional[str] = None
+    attributes: List[EntityAttribute] = []
+
+
+class EventEntityList(BaseModel):
+    events: List[EventWithEntities]

cognee/temporal_poc/temporal_cognify.py

@@ -0,0 +1,64 @@
+from typing import Union, Optional, List
+from uuid import UUID
+from cognee.shared.logging_utils import get_logger
+from cognee.shared.data_models import KnowledgeGraph
+from cognee.infrastructure.llm import get_max_chunk_tokens, get_llm_config
+
+from cognee.api.v1.cognify.cognify import run_cognify_blocking
+from cognee.modules.pipelines.tasks.task import Task
+from cognee.modules.chunking.TextChunker import TextChunker
+from cognee.modules.ontology.rdf_xml.OntologyResolver import OntologyResolver
+
+from cognee.modules.users.models import User
+from cognee.tasks.documents import (
+    check_permissions_on_dataset,
+    classify_documents,
+    extract_chunks_from_documents,
+)
+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.temporal_poc.event_extraction import extract_events_and_entities
+from cognee.temporal_poc.event_knowledge_graph import process_event_knowledge_graph
+
+logger = get_logger("temporal_cognify")
+
+
+async def get_temporal_tasks(
+    user: User = None, chunker=TextChunker, chunk_size: int = None
+) -> list[Task]:
+    temporal_tasks = [
+        Task(classify_documents),
+        Task(check_permissions_on_dataset, user=user, permissions=["write"]),
+        Task(
+            extract_chunks_from_documents,
+            max_chunk_size=chunk_size or get_max_chunk_tokens(),
+            chunker=chunker,
+        ),
+        Task(extract_events_and_entities, task_config={"chunk_size": 10}),
+        Task(process_event_knowledge_graph),
+        Task(add_data_points, task_config={"batch_size": 10}),
+    ]
+
+    return temporal_tasks
+
+
+async def temporal_cognify(
+    datasets: Union[str, list[str], list[UUID]] = None,
+    user: User = None,
+    chunker=TextChunker,
+    chunk_size: int = None,
+    vector_db_config: dict = None,
+    graph_db_config: dict = None,
+    incremental_loading: bool = True,
+):
+    tasks = await get_temporal_tasks(user, chunker, chunk_size)
+
+    return await run_cognify_blocking(
+        tasks=tasks,
+        user=user,
+        datasets=datasets,
+        vector_db_config=vector_db_config,
+        graph_db_config=graph_db_config,
+        incremental_loading=incremental_loading,
+    )

cognee/temporal_poc/temporal_example.py

@@ -0,0 +1,48 @@
+import asyncio
+import cognee
+from cognee.shared.logging_utils import setup_logging, INFO
+from cognee.temporal_poc.temporal_cognify import temporal_cognify
+
+import json
+from pathlib import Path
+
+
+async def reading_temporal_data():
+    path = Path("cognee/temporal_poc/test_hard.json")
+    contexts = []
+    seen = set()
+    with path.open(encoding="utf-8") as f:
+        for line in f:
+            line = line.strip()
+            if not line:
+                continue
+            entry = json.loads(line)
+            ctx = entry.get("context", "")
+            if ctx and ctx not in seen:
+                seen.add(ctx)
+                contexts.append(ctx)
+    return contexts
+
+
+async def main():
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+
+    texts = await reading_temporal_data()
+    texts = texts[:5]
+
+    await cognee.add(texts)
+    await temporal_cognify()
+
+    print()
+
+
+if __name__ == "__main__":
+    logger = setup_logging(log_level=INFO)
+
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+    try:
+        loop.run_until_complete(main())
+    finally:
+        loop.run_until_complete(loop.shutdown_asyncgens())