cognee 0.2.3.dev1__py3-none-any.whl → 0.3.0.dev0__py3-none-any.whl

cognee/modules/engine/utils/generate_event_datapoint.py

@@ -0,0 +1,46 @@
+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:
+    """
+    Generates an Event datapoint from a given event model, including temporal metadata if available.
+
+    The function maps the basic attributes (name, description, location) from the input event
+    and enriches them with temporal information. If start and end times are provided, an
+    Interval is created. If only one timestamp is available, it is added directly. Temporal
+    information is also appended to the event description for context.
+
+    Args:
+        event: An event model instance containing attributes such as name, description,
+               location, time_from, and time_to.
+
+    Returns:
+        Event: A structured Event object with name, description, location, and enriched
+               temporal details.
+    """
+    # Base event data
+    event_data = {
+        "name": event.name,
+        "description": event.description,
+        "location": event.location,
+    }
+
+    # Create timestamps if they exist
+    time_from = generate_timestamp_datapoint(event.time_from) if event.time_from else None
+    time_to = generate_timestamp_datapoint(event.time_to) if event.time_to else None
+
+    # Add temporal information
+    if time_from and time_to:
+        event_data["during"] = Interval(time_from=time_from, time_to=time_to)
+        # Enrich description with temporal info
+        temporal_info = f"\n---\nTime data: {time_from.timestamp_str} to {time_to.timestamp_str}"
+        event_data["description"] = (event_data["description"] or "Event") + temporal_info
+    elif time_from or time_to:
+        timestamp = time_from or time_to
+        event_data["at"] = timestamp
+        # Enrich description with temporal info
+        temporal_info = f"\n---\nTime data: {timestamp.timestamp_str}"
+        event_data["description"] = (event_data["description"] or "Event") + temporal_info
+
+    return Event(**event_data)

cognee/modules/engine/utils/generate_timestamp_datapoint.py

@@ -0,0 +1,51 @@
+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:
+    """
+    Generates a normalized Timestamp datapoint from a given Timestamp model.
+
+    The function converts the provided timestamp into an integer representation,
+    constructs a human-readable string format, and creates a new Timestamp object
+    with a unique identifier.
+
+    Args:
+        ts (Timestamp): The input Timestamp model containing date and time components.
+
+    Returns:
+        Timestamp: A new Timestamp object with a generated ID, integer representation,
+                   original components, and formatted string.
+    """
+
+    time_at = date_to_int(ts)
+    timestamp_str = (
+        f"{ts.year:04d}-{ts.month:02d}-{ts.day:02d} {ts.hour:02d}:{ts.minute:02d}:{ts.second:02d}"
+    )
+    return Timestamp(
+        id=generate_node_id(str(time_at)),
+        time_at=time_at,
+        year=ts.year,
+        month=ts.month,
+        day=ts.day,
+        hour=ts.hour,
+        minute=ts.minute,
+        second=ts.second,
+        timestamp_str=timestamp_str,
+    )
+
+
+def date_to_int(ts: Timestamp) -> int:
+    """
+    Converts a Timestamp model into an integer representation in milliseconds since the Unix epoch (UTC).
+
+    Args:
+        ts (Timestamp): The input Timestamp model containing year, month, day, hour, minute, and second.
+
+    Returns:
+        int: The UTC timestamp in milliseconds since January 1, 1970.
+    """
+    dt = datetime(ts.year, ts.month, ts.day, ts.hour, ts.minute, ts.second, tzinfo=timezone.utc)
+    time = int(dt.timestamp() * 1000)
+    return time

cognee/modules/graph/cognee_graph/CogneeGraph.py

@@ -76,7 +76,7 @@ class CogneeGraph(CogneeAbstractGraph):
             start_time = time.time()
 
             # Determine projection strategy
-            if node_type is not None and node_name is not None:
+            if node_type is not None and node_name not in [None, [], ""]:
                 nodes_data, edges_data = await adapter.get_nodeset_subgraph(
                     node_type=node_type, node_name=node_name
                 )
@@ -180,7 +180,7 @@ class CogneeGraph(CogneeAbstractGraph):
             logger.error(f"Error mapping vector distances to edges: {str(ex)}")
             raise ex
 
-    async def calculate_top_triplet_importances(self, k: int) -> List:
+    async def calculate_top_triplet_importances(self, k: int) -> List[Edge]:
         def score(edge):
             n1 = edge.node1.attributes.get("vector_distance", 1)
             n2 = edge.node2.attributes.get("vector_distance", 1)

cognee/modules/graph/methods/get_formatted_graph_data.py

@@ -3,10 +3,11 @@ from cognee.infrastructure.databases.graph import get_graph_engine
 from cognee.context_global_variables import set_database_global_context_variables
 from cognee.modules.data.exceptions.exceptions import DatasetNotFoundError
 from cognee.modules.data.methods import get_authorized_dataset
+from cognee.modules.users.models import User
 
 
-async def get_formatted_graph_data(dataset_id: UUID, user_id: UUID):
-    dataset = await get_authorized_dataset(user_id, dataset_id)
+async def get_formatted_graph_data(dataset_id: UUID, user: User):
+    dataset = await get_authorized_dataset(user, dataset_id)
     if not dataset:
         raise DatasetNotFoundError(message="Dataset not found.")
 

cognee/modules/graph/utils/__init__.py

@@ -4,3 +4,4 @@ from .get_model_instance_from_graph import get_model_instance_from_graph
 from .retrieve_existing_edges import retrieve_existing_edges
 from .convert_node_to_data_point import convert_node_to_data_point
 from .deduplicate_nodes_and_edges import deduplicate_nodes_and_edges
+from .resolve_edges_to_text import resolve_edges_to_text

cognee/modules/graph/utils/resolve_edges_to_text.py

@@ -0,0 +1,71 @@
+from typing import List
+from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge
+
+
+async def resolve_edges_to_text(retrieved_edges: List[Edge]) -> str:
+    """
+    Converts retrieved graph edges into a human-readable string format.
+
+    Parameters:
+    -----------
+
+        - retrieved_edges (list): A list of edges retrieved from the graph.
+
+    Returns:
+    --------
+
+        - str: A formatted string representation of the nodes and their connections.
+    """
+
+    def _get_nodes(retrieved_edges: List[Edge]) -> dict:
+        def _get_title(text: str, first_n_words: int = 7, top_n_words: int = 3) -> str:
+            def _top_n_words(text, stop_words=None, top_n=3, separator=", "):
+                """Concatenates the top N frequent words in text."""
+                if stop_words is None:
+                    from cognee.modules.retrieval.utils.stop_words import DEFAULT_STOP_WORDS
+
+                    stop_words = DEFAULT_STOP_WORDS
+
+                import string
+
+                words = [word.lower().strip(string.punctuation) for word in text.split()]
+
+                if stop_words:
+                    words = [word for word in words if word and word not in stop_words]
+
+                from collections import Counter
+
+                top_words = [word for word, freq in Counter(words).most_common(top_n)]
+
+                return separator.join(top_words)
+
+            """Creates a title, by combining first words with most frequent words from the text."""
+            first_words = text.split()[:first_n_words]
+            top_words = _top_n_words(text, top_n=first_n_words)
+            return f"{' '.join(first_words)}... [{top_words}]"
+
+        """Creates a dictionary of nodes with their names and content."""
+        nodes = {}
+        for edge in retrieved_edges:
+            for node in (edge.node1, edge.node2):
+                if node.id not in nodes:
+                    text = node.attributes.get("text")
+                    if text:
+                        name = _get_title(text)
+                        content = text
+                    else:
+                        name = node.attributes.get("name", "Unnamed Node")
+                        content = node.attributes.get("description", name)
+                    nodes[node.id] = {"node": node, "name": name, "content": content}
+        return nodes
+
+    nodes = _get_nodes(retrieved_edges)
+    node_section = "\n".join(
+        f"Node: {info['name']}\n__node_content_start__\n{info['content']}\n__node_content_end__\n"
+        for info in nodes.values()
+    )
+    connection_section = "\n".join(
+        f"{nodes[edge.node1.id]['name']} --[{edge.attributes['relationship_type']}]--> {nodes[edge.node2.id]['name']}"
+        for edge in retrieved_edges
+    )
+    return f"Nodes:\n{node_section}\n\nConnections:\n{connection_section}"

cognee/modules/memify/__init__.py

@@ -0,0 +1 @@
+from .memify import memify

cognee/modules/memify/memify.py

@@ -0,0 +1,118 @@
+from typing import Union, Optional, List, Type, Any
+from uuid import UUID
+
+from cognee.shared.logging_utils import get_logger
+
+from cognee.modules.retrieval.utils.brute_force_triplet_search import get_memory_fragment
+from cognee.context_global_variables import set_database_global_context_variables
+from cognee.modules.engine.models.node_set import NodeSet
+from cognee.modules.pipelines import run_pipeline
+from cognee.modules.pipelines.tasks.task import Task
+from cognee.modules.users.models import User
+from cognee.modules.pipelines.layers.resolve_authorized_user_datasets import (
+    resolve_authorized_user_datasets,
+)
+from cognee.modules.pipelines.layers.reset_dataset_pipeline_run_status import (
+    reset_dataset_pipeline_run_status,
+)
+from cognee.modules.engine.operations.setup import setup
+from cognee.modules.pipelines.layers.pipeline_execution_mode import get_pipeline_executor
+from cognee.tasks.memify.extract_subgraph_chunks import extract_subgraph_chunks
+from cognee.tasks.codingagents.coding_rule_associations import (
+    add_rule_associations,
+)
+
+logger = get_logger("memify")
+
+
+async def memify(
+    extraction_tasks: Union[List[Task], List[str]] = None,
+    enrichment_tasks: Union[List[Task], List[str]] = None,
+    data: Optional[Any] = None,
+    dataset: Union[str, UUID] = "main_dataset",
+    user: User = None,
+    node_type: Optional[Type] = NodeSet,
+    node_name: Optional[List[str]] = None,
+    vector_db_config: Optional[dict] = None,
+    graph_db_config: Optional[dict] = None,
+    run_in_background: bool = False,
+):
+    """
+    Enrichment pipeline in Cognee, can work with already built graphs. If no data is provided existing knowledge graph will be used as data,
+    custom data can also be provided instead which can be processed with provided extraction and enrichment tasks.
+
+    Provided tasks and data will be arranged to run the Cognee pipeline and execute graph enrichment/creation.
+
+    This is the core processing step in Cognee that converts raw text and documents
+    into an intelligent knowledge graph. It analyzes content, extracts entities and
+    relationships, and creates semantic connections for enhanced search and reasoning.
+
+    Args:
+        extraction_tasks: List of Cognee Tasks to execute for graph/data extraction.
+        enrichment_tasks: List of Cognee Tasks to handle enrichment of provided graph/data from extraction tasks.
+        data: The data to ingest. Can be anything when custom extraction and enrichment tasks are used.
+              Data provided here will be forwarded to the first extraction task in the pipeline as input.
+              If no data is provided the whole graph (or subgraph if node_name/node_type is specified) will be forwarded
+        dataset: Dataset name or dataset uuid to process.
+        user: User context for authentication and data access. Uses default if None.
+        node_type: Filter graph to specific entity types (for advanced filtering). Used when no data is provided.
+        node_name: Filter graph to specific named entities (for targeted search). Used when no data is provided.
+        vector_db_config: Custom vector database configuration for embeddings storage.
+        graph_db_config: Custom graph database configuration for relationship storage.
+        run_in_background: If True, starts processing asynchronously and returns immediately.
+                          If False, waits for completion before returning.
+                          Background mode recommended for large datasets (>100MB).
+                          Use pipeline_run_id from return value to monitor progress.
+    """
+
+    # Use default coding rules tasks if no tasks were provided
+    if not extraction_tasks:
+        extraction_tasks = [Task(extract_subgraph_chunks)]
+    if not enrichment_tasks:
+        enrichment_tasks = [
+            Task(
+                add_rule_associations,
+                rules_nodeset_name="coding_agent_rules",
+                task_config={"batch_size": 1},
+            )
+        ]
+
+    await setup()
+
+    user, authorized_dataset_list = await resolve_authorized_user_datasets(dataset, user)
+    authorized_dataset = authorized_dataset_list[0]
+
+    if not data:
+        # Will only be used if ENABLE_BACKEND_ACCESS_CONTROL is set to True
+        await set_database_global_context_variables(
+            authorized_dataset.id, authorized_dataset.owner_id
+        )
+
+        memory_fragment = await get_memory_fragment(node_type=node_type, node_name=node_name)
+        # Subgraphs should be a single element in the list to represent one data item
+        data = [memory_fragment]
+
+    memify_tasks = [
+        *extraction_tasks,  # Unpack tasks provided to memify pipeline
+        *enrichment_tasks,
+    ]
+
+    await reset_dataset_pipeline_run_status(
+        authorized_dataset.id, user, pipeline_names=["memify_pipeline"]
+    )
+
+    # By calling get pipeline executor we get a function that will have the run_pipeline run in the background or a function that we will need to wait for
+    pipeline_executor_func = get_pipeline_executor(run_in_background=run_in_background)
+
+    # Run the run_pipeline in the background or blocking based on executor
+    return await pipeline_executor_func(
+        pipeline=run_pipeline,
+        tasks=memify_tasks,
+        user=user,
+        data=data,
+        datasets=authorized_dataset.id,
+        vector_db_config=vector_db_config,
+        graph_db_config=graph_db_config,
+        incremental_loading=False,
+        pipeline_name="memify_pipeline",
+    )

cognee/modules/notebooks/methods/__init__.py

@@ -0,0 +1,5 @@
+from .get_notebook import get_notebook
+from .get_notebooks import get_notebooks
+from .create_notebook import create_notebook
+from .update_notebook import update_notebook
+from .delete_notebook import delete_notebook

cognee/modules/notebooks/methods/create_notebook.py

@@ -0,0 +1,26 @@
+from uuid import UUID
+from typing import List, Optional
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from cognee.infrastructure.databases.relational import with_async_session
+
+from ..models.Notebook import Notebook, NotebookCell
+
+
+@with_async_session
+async def create_notebook(
+    user_id: UUID,
+    notebook_name: str,
+    cells: Optional[List[NotebookCell]],
+    deletable: Optional[bool],
+    session: AsyncSession,
+) -> Notebook:
+    notebook = Notebook(
+        name=notebook_name, owner_id=user_id, cells=cells, deletable=deletable or True
+    )
+
+    session.add(notebook)
+
+    await session.commit()
+
+    return notebook

cognee/modules/notebooks/methods/delete_notebook.py

@@ -0,0 +1,13 @@
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from cognee.infrastructure.databases.relational import with_async_session
+
+from ..models.Notebook import Notebook
+
+
+@with_async_session
+async def delete_notebook(
+    notebook: Notebook,
+    session: AsyncSession,
+) -> None:
+    await session.delete(notebook)

cognee/modules/notebooks/methods/get_notebook.py

@@ -0,0 +1,21 @@
+from uuid import UUID
+from typing import Optional
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from cognee.infrastructure.databases.relational import with_async_session
+
+from ..models.Notebook import Notebook
+
+
+@with_async_session
+async def get_notebook(
+    notebook_id: UUID,
+    user_id: UUID,
+    session: AsyncSession,
+) -> Optional[Notebook]:
+    result = await session.execute(
+        select(Notebook).where(Notebook.owner_id == user_id and Notebook.id == notebook_id)
+    )
+
+    return result.scalar()

cognee/modules/notebooks/methods/get_notebooks.py

@@ -0,0 +1,18 @@
+from uuid import UUID
+from typing import List
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from cognee.infrastructure.databases.relational import with_async_session
+
+from ..models.Notebook import Notebook
+
+
+@with_async_session
+async def get_notebooks(
+    user_id: UUID,
+    session: AsyncSession,
+) -> List[Notebook]:
+    result = await session.execute(select(Notebook).where(Notebook.owner_id == user_id))
+
+    return list(result.scalars().all())

cognee/modules/notebooks/methods/update_notebook.py

@@ -0,0 +1,17 @@
+from typing import Callable, AsyncContextManager
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from cognee.infrastructure.databases.relational import with_async_session
+
+from ..models.Notebook import Notebook
+
+
+@with_async_session
+async def update_notebook(
+    notebook: Notebook,
+    session: AsyncSession,
+) -> Notebook:
+    if notebook not in session:
+        session.add(notebook)
+
+    return notebook

cognee/modules/notebooks/models/Notebook.py

@@ -0,0 +1,53 @@
+import json
+from typing import List, Literal
+from uuid import uuid4, UUID as UUID_t
+from pydantic import BaseModel, ConfigDict
+from datetime import datetime, timezone
+from fastapi.encoders import jsonable_encoder
+from sqlalchemy import Boolean, Column, DateTime, JSON, UUID, String, TypeDecorator
+from sqlalchemy.orm import mapped_column, Mapped
+
+from cognee.infrastructure.databases.relational import Base
+
+
+class NotebookCell(BaseModel):
+    id: UUID_t
+    type: Literal["markdown", "code"]
+    name: str
+    content: str
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+
+class NotebookCellList(TypeDecorator):
+    impl = JSON
+    cache_ok = True
+
+    def process_bind_param(self, notebook_cells, dialect):
+        if notebook_cells is None:
+            return []
+        return [
+            json.dumps(jsonable_encoder(cell)) if isinstance(cell, NotebookCell) else cell
+            for cell in notebook_cells
+        ]
+
+    def process_result_value(self, cells_json_list, dialect):
+        if cells_json_list is None:
+            return []
+        return [NotebookCell(**json.loads(json_string)) for json_string in cells_json_list]
+
+
+class Notebook(Base):
+    __tablename__ = "notebooks"
+
+    id: Mapped[UUID_t] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid4)
+
+    owner_id: Mapped[UUID_t] = mapped_column(UUID(as_uuid=True), index=True)
+
+    name: Mapped[str] = mapped_column(String, nullable=False)
+
+    cells: Mapped[List[NotebookCell]] = mapped_column(NotebookCellList, nullable=False)
+
+    deletable: Mapped[bool] = mapped_column(Boolean, default=True)
+
+    created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))

cognee/modules/notebooks/models/__init__.py

@@ -0,0 +1 @@
+from .Notebook import Notebook, NotebookCell

cognee/modules/notebooks/operations/__init__.py

@@ -0,0 +1 @@
+from .run_in_local_sandbox import run_in_local_sandbox

cognee/modules/notebooks/operations/run_in_local_sandbox.py

@@ -0,0 +1,55 @@
+import io
+import sys
+import traceback
+
+
+def wrap_in_async_handler(user_code: str) -> str:
+    return (
+        "from cognee.infrastructure.utils.run_sync import run_sync\n\n"
+        "async def __user_main__():\n"
+        + "\n".join("    " + line for line in user_code.strip().split("\n"))
+        + "\n"
+        "    globals().update(locals())\n\n"
+        "run_sync(__user_main__())\n"
+    )
+
+
+def run_in_local_sandbox(code, environment=None):
+    environment = environment or {}
+    code = wrap_in_async_handler(code.replace("\xa0", "\n"))
+
+    buffer = io.StringIO()
+    sys_stdout = sys.stdout
+    sys.stdout = buffer
+    sys.stderr = buffer
+
+    error = None
+
+    printOutput = []
+
+    def customPrintFunction(output):
+        printOutput.append(output)
+
+    environment["print"] = customPrintFunction
+
+    try:
+        exec(code, environment)
+    except Exception:
+        error = traceback.format_exc()
+    finally:
+        sys.stdout = sys_stdout
+        sys.stderr = sys_stdout
+
+    return printOutput, error
+
+
+if __name__ == "__main__":
+    run_in_local_sandbox("""
+import cognee
+
+await cognee.add("Test file with some random content 3.")
+
+a = "asd"
+
+b = {"c": "dfgh"}
+""")

cognee/modules/pipelines/__init__.py

@@ -1,4 +1,4 @@
 from .tasks.task import Task
 from .operations.run_tasks import run_tasks
 from .operations.run_parallel import run_tasks_parallel
-from .operations.pipeline import cognee_pipeline
+from .operations.pipeline import run_pipeline

cognee/modules/pipelines/exceptions/tasks.py

@@ -0,0 +1,18 @@
+from fastapi import status
+from cognee.exceptions import CogneeValidationError
+
+
+class WrongTaskTypeError(CogneeValidationError):
+    """
+    Raised when the tasks argument is not a list of Task class instances.
+    """
+
+    def __init__(
+        self,
+        message: str = "tasks argument must be a list, containing Task class instances.",
+        name: str = "WrongTaskTypeError",
+        status_code=status.HTTP_400_BAD_REQUEST,
+    ):
+        self.message = message
+        self.name = name
+        self.status_code = status_code

cognee/modules/pipelines/layers/__init__.py

@@ -0,0 +1 @@
+from .validate_pipeline_tasks import validate_pipeline_tasks

cognee/modules/pipelines/layers/check_pipeline_run_qualification.py

@@ -0,0 +1,59 @@
+from typing import Union, Optional
+from cognee.modules.data.models import Dataset
+from cognee.modules.data.models import Data
+from cognee.modules.pipelines.models import PipelineRunStatus
+from cognee.modules.pipelines.operations.get_pipeline_status import get_pipeline_status
+from cognee.modules.pipelines.methods import get_pipeline_run_by_dataset
+from cognee.shared.logging_utils import get_logger
+
+from cognee.modules.pipelines.models.PipelineRunInfo import (
+    PipelineRunCompleted,
+    PipelineRunStarted,
+)
+
+logger = get_logger(__name__)
+
+
+async def check_pipeline_run_qualification(
+    dataset: Dataset, data: list[Data], pipeline_name: str
+) -> Optional[Union[PipelineRunStarted, PipelineRunCompleted]]:
+    """
+    Function used to determine if pipeline is currently being processed or was already processed.
+    In case pipeline was or is being processed return value is returned and current pipline execution should be stopped.
+    In case pipeline is not or was not processed there will be no return value and pipeline processing can start.
+
+    Args:
+        dataset: Dataset object
+        data: List of Data
+        pipeline_name: pipeline name
+
+    Returns: Pipeline state if it is being processed or was already processed
+
+    """
+
+    # async with update_status_lock: TODO: Add UI lock to prevent multiple backend requests
+    if isinstance(dataset, Dataset):
+        task_status = await get_pipeline_status([dataset.id], pipeline_name)
+    else:
+        task_status = {}
+
+    if str(dataset.id) in task_status:
+        if task_status[str(dataset.id)] == PipelineRunStatus.DATASET_PROCESSING_STARTED:
+            logger.info("Dataset %s is already being processed.", dataset.id)
+            pipeline_run = await get_pipeline_run_by_dataset(dataset.id, pipeline_name)
+            return PipelineRunStarted(
+                pipeline_run_id=pipeline_run.pipeline_run_id,
+                dataset_id=dataset.id,
+                dataset_name=dataset.name,
+                payload=data,
+            )
+        elif task_status[str(dataset.id)] == PipelineRunStatus.DATASET_PROCESSING_COMPLETED:
+            logger.info("Dataset %s is already processed.", dataset.id)
+            pipeline_run = await get_pipeline_run_by_dataset(dataset.id, pipeline_name)
+            return PipelineRunCompleted(
+                pipeline_run_id=pipeline_run.pipeline_run_id,
+                dataset_id=dataset.id,
+                dataset_name=dataset.name,
+            )
+
+    return