cognee 0.5.1.dev0__py3-none-any.whl → 0.5.2__py3-none-any.whl

cognee/modules/graph/cognee_graph/CogneeGraph.py

@@ -1,6 +1,7 @@
 import time
 from cognee.shared.logging_utils import get_logger
-from typing import List, Dict, Union, Optional, Type
+from cognee.modules.engine.utils.generate_edge_id import generate_edge_id
+from typing import List, Dict, Union, Optional, Type, Iterable, Tuple, Callable, Any
 
 from cognee.modules.graph.exceptions import (
     EntityNotFoundError,
@@ -25,12 +26,16 @@ class CogneeGraph(CogneeAbstractGraph):
 
     nodes: Dict[str, Node]
     edges: List[Edge]
+    edges_by_distance_key: Dict[str, List[Edge]]
     directed: bool
+    triplet_distance_penalty: float
 
     def __init__(self, directed: bool = True):
         self.nodes = {}
         self.edges = []
+        self.edges_by_distance_key = {}
         self.directed = directed
+        self.triplet_distance_penalty = 3.5
 
     def add_node(self, node: Node) -> None:
         if node.id not in self.nodes:
@@ -40,8 +45,20 @@ class CogneeGraph(CogneeAbstractGraph):
 
     def add_edge(self, edge: Edge) -> None:
         self.edges.append(edge)
+
+        edge_text = edge.attributes.get("edge_text") or edge.attributes.get("relationship_type")
+        edge.attributes["edge_type_id"] = (
+            generate_edge_id(edge_id=edge_text) if edge_text else None
+        )  # Update edge with generated edge_type_id
+
         edge.node1.add_skeleton_edge(edge)
         edge.node2.add_skeleton_edge(edge)
+        key = edge.get_distance_key()
+        if not key:
+            return
+        if key not in self.edges_by_distance_key:
+            self.edges_by_distance_key[key] = []
+        self.edges_by_distance_key[key].append(edge)
 
     def get_node(self, node_id: str) -> Node:
         return self.nodes.get(node_id, None)
@@ -56,6 +73,29 @@ class CogneeGraph(CogneeAbstractGraph):
     def get_edges(self) -> List[Edge]:
         return self.edges
 
+    def reset_distances(self, collection: Iterable[Union[Node, Edge]], query_count: int) -> None:
+        """Reset vector distances for a collection of nodes or edges."""
+        for item in collection:
+            item.reset_vector_distances(query_count, self.triplet_distance_penalty)
+
+    def _normalize_query_distance_lists(
+        self, distances: List, query_list_length: Optional[int] = None, name: str = "distances"
+    ) -> List:
+        """Normalize shape: flat list -> single-query; nested list -> multi-query."""
+        if not distances:
+            return []
+        first_item = distances[0]
+        if isinstance(first_item, (list, tuple)):
+            per_query_lists = distances
+        else:
+            per_query_lists = [distances]
+        if query_list_length is not None and len(per_query_lists) != query_list_length:
+            raise ValueError(
+                f"{name} has {len(per_query_lists)} query lists, "
+                f"but query_list_length is {query_list_length}"
+            )
+        return per_query_lists
+
     async def _get_nodeset_subgraph(
         self,
         adapter,
@@ -148,7 +188,7 @@ class CogneeGraph(CogneeAbstractGraph):
                     adapter, memory_fragment_filter
                 )
 
-            import time
+            self.triplet_distance_penalty = triplet_distance_penalty
 
             start_time = time.time()
             # Process nodes
@@ -182,9 +222,6 @@ class CogneeGraph(CogneeAbstractGraph):
                         edge_penalty=triplet_distance_penalty,
                     )
                     self.add_edge(edge)
-
-                    source_node.add_skeleton_edge(edge)
-                    target_node.add_skeleton_edge(edge)
                 else:
                     raise EntityNotFoundError(
                         message=f"Edge references nonexistent nodes: {source_id} -> {target_id}"
@@ -200,41 +237,117 @@ class CogneeGraph(CogneeAbstractGraph):
             logger.error(f"Error during graph projection: {str(e)}")
             raise
 
-    async def map_vector_distances_to_graph_nodes(self, node_distances) -> None:
-        mapped_nodes = 0
-        for category, scored_results in node_distances.items():
-            for scored_result in scored_results:
-                node_id = str(scored_result.id)
-                score = scored_result.score
-                node = self.get_node(node_id)
-                if node:
-                    node.add_attribute("vector_distance", score)
-                    mapped_nodes += 1
-
-    async def map_vector_distances_to_graph_edges(self, edge_distances) -> None:
-        try:
-            if edge_distances is None:
-                return
+    async def map_vector_distances_to_graph_nodes(
+        self,
+        node_distances,
+        query_list_length: Optional[int] = None,
+    ) -> None:
+        """Map vector distances to nodes, supporting single- and multi-query input shapes."""
 
-            embedding_map = {result.payload["text"]: result.score for result in edge_distances}
+        query_count = query_list_length or 1
 
-            for edge in self.edges:
-                edge_key = edge.attributes.get("edge_text") or edge.attributes.get(
-                    "relationship_type"
-                )
-                distance = embedding_map.get(edge_key, None)
-                if distance is not None:
-                    edge.attributes["vector_distance"] = distance
+        self.reset_distances(self.nodes.values(), query_count)
+
+        for collection_name, scored_results in node_distances.items():
+            if not scored_results:
+                continue
+
+            per_query_scored_results = self._normalize_query_distance_lists(
+                scored_results, query_list_length, f"Collection '{collection_name}'"
+            )
+
+            for query_index, scored_results in enumerate(per_query_scored_results):
+                for result in scored_results:
+                    node_id = str(getattr(result, "id", None))
+                    if not node_id:
+                        continue
+                    node = self.get_node(node_id)
+                    if node is None:
+                        continue
+                    score = float(getattr(result, "score", self.triplet_distance_penalty))
+                    node.update_distance_for_query(
+                        query_index=query_index,
+                        score=score,
+                        query_count=query_count,
+                        default_penalty=self.triplet_distance_penalty,
+                    )
 
-        except Exception as ex:
-            logger.error(f"Error mapping vector distances to edges: {str(ex)}")
-            raise ex
+    async def map_vector_distances_to_graph_edges(
+        self,
+        edge_distances,
+        query_list_length: Optional[int] = None,
+    ) -> None:
+        """Map vector distances to graph edges, supporting single- and multi-query input shapes."""
+        query_count = query_list_length or 1
+
+        self.reset_distances(self.edges, query_count)
+
+        if not edge_distances:
+            return None
 
-    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)
-            e = edge.attributes.get("vector_distance", 1)
-            return n1 + n2 + e
+        per_query_scored_results = self._normalize_query_distance_lists(
+            edge_distances, query_list_length, "edge_distances"
+        )
+
+        for query_index, scored_results in enumerate(per_query_scored_results):
+            for result in scored_results:
+                matching_edges = self.edges_by_distance_key.get(str(result.id))
+                if not matching_edges:
+                    continue
+                for edge in matching_edges:
+                    edge.update_distance_for_query(
+                        query_index=query_index,
+                        score=float(getattr(result, "score", self.triplet_distance_penalty)),
+                        query_count=query_count,
+                        default_penalty=self.triplet_distance_penalty,
+                    )
+
+    def _calculate_query_top_triplet_importances(
+        self,
+        k: int,
+        query_index: int = 0,
+    ) -> List[Edge]:
+        """Calculate top k triplet importances for a specific query index."""
+
+        def score(edge: Edge) -> float:
+            elements = (
+                (edge.node1, f"node {edge.node1.id}"),
+                (edge.node2, f"node {edge.node2.id}"),
+                (edge, f"edge {edge.node1.id}->{edge.node2.id}"),
+            )
+
+            importances = []
+            for element, label in elements:
+                distances = element.attributes.get("vector_distance")
+                if not isinstance(distances, list) or query_index >= len(distances):
+                    raise ValueError(
+                        f"{label}: vector_distance must be a list with length > {query_index} "
+                        f"before scoring (got {type(distances).__name__} with length "
+                        f"{len(distances) if isinstance(distances, list) else 'n/a'})"
+                    )
+                value = distances[query_index]
+                try:
+                    importances.append(float(value))
+                except (TypeError, ValueError):
+                    raise ValueError(
+                        f"{label}: vector_distance[{query_index}] must be float-like, "
+                        f"got {type(value).__name__}"
+                    )
+
+            return sum(importances)
 
         return heapq.nsmallest(k, self.edges, key=score)
+
+    async def calculate_top_triplet_importances(
+        self, k: int, query_list_length: Optional[int] = None
+    ) -> Union[List[Edge], List[List[Edge]]]:
+        """Calculate top k triplet importances, supporting both single and multi-query modes."""
+        query_count = query_list_length or 1
+        results = [
+            self._calculate_query_top_triplet_importances(k=k, query_index=i)
+            for i in range(query_count)
+        ]
+
+        if query_list_length is None:
+            return results[0]
+        return results

cognee/modules/graph/cognee_graph/CogneeGraphElements.py

@@ -30,11 +30,31 @@ class Node:
             raise InvalidDimensionsError()
         self.id = node_id
         self.attributes = attributes if attributes is not None else {}
-        self.attributes["vector_distance"] = node_penalty
+        self.attributes["vector_distance"] = None
         self.skeleton_neighbours = []
         self.skeleton_edges = []
         self.status = np.ones(dimension, dtype=int)
 
+    def reset_vector_distances(self, query_count: int, default_penalty: float) -> None:
+        self.attributes["vector_distance"] = [default_penalty] * query_count
+
+    def ensure_vector_distance_list(self, query_count: int, default_penalty: float) -> List[float]:
+        distances = self.attributes.get("vector_distance")
+        if not isinstance(distances, list) or len(distances) != query_count:
+            distances = [default_penalty] * query_count
+            self.attributes["vector_distance"] = distances
+        return distances
+
+    def update_distance_for_query(
+        self,
+        query_index: int,
+        score: float,
+        query_count: int,
+        default_penalty: float,
+    ) -> None:
+        distances = self.ensure_vector_distance_list(query_count, default_penalty)
+        distances[query_index] = score
+
     def add_skeleton_neighbor(self, neighbor: "Node") -> None:
         if neighbor not in self.skeleton_neighbours:
             self.skeleton_neighbours.append(neighbor)
@@ -116,10 +136,36 @@ class Edge:
         self.node1 = node1
         self.node2 = node2
         self.attributes = attributes if attributes is not None else {}
-        self.attributes["vector_distance"] = edge_penalty
+        self.attributes["vector_distance"] = None
         self.directed = directed
         self.status = np.ones(dimension, dtype=int)
 
+    def get_distance_key(self) -> Optional[str]:
+        key = self.attributes.get("edge_type_id")
+        if key is None:
+            return None
+        return str(key)
+
+    def reset_vector_distances(self, query_count: int, default_penalty: float) -> None:
+        self.attributes["vector_distance"] = [default_penalty] * query_count
+
+    def ensure_vector_distance_list(self, query_count: int, default_penalty: float) -> List[float]:
+        distances = self.attributes.get("vector_distance")
+        if not isinstance(distances, list) or len(distances) != query_count:
+            distances = [default_penalty] * query_count
+            self.attributes["vector_distance"] = distances
+        return distances
+
+    def update_distance_for_query(
+        self,
+        query_index: int,
+        score: float,
+        query_count: int,
+        default_penalty: float,
+    ) -> None:
+        distances = self.ensure_vector_distance_list(query_count, default_penalty)
+        distances[query_index] = score
+
     def is_edge_alive_in_dimension(self, dimension: int) -> bool:
         if dimension < 0 or dimension >= len(self.status):
             raise DimensionOutOfRangeError(dimension=dimension, max_index=len(self.status) - 1)

cognee/modules/graph/utils/__init__.py

@@ -5,3 +5,4 @@ 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
+from .get_entity_nodes_from_triplets import get_entity_nodes_from_triplets

cognee/modules/graph/utils/get_entity_nodes_from_triplets.py

@@ -0,0 +1,12 @@
+def get_entity_nodes_from_triplets(triplets):
+    entity_nodes = []
+    seen_ids = set()
+    for triplet in triplets:
+        if hasattr(triplet, "node1") and triplet.node1 and triplet.node1.id not in seen_ids:
+            entity_nodes.append({"id": str(triplet.node1.id)})
+            seen_ids.add(triplet.node1.id)
+        if hasattr(triplet, "node2") and triplet.node2 and triplet.node2.id not in seen_ids:
+            entity_nodes.append({"id": str(triplet.node2.id)})
+            seen_ids.add(triplet.node2.id)
+
+    return entity_nodes

cognee/modules/notebooks/methods/__init__.py

@@ -3,3 +3,4 @@ from .get_notebooks import get_notebooks
 from .create_notebook import create_notebook
 from .update_notebook import update_notebook
 from .delete_notebook import delete_notebook
+from .create_tutorial_notebooks import create_tutorial_notebooks

cognee/modules/notebooks/methods/create_notebook.py

@@ -6,40 +6,6 @@ from cognee.infrastructure.databases.relational import with_async_session
 
 from ..models.Notebook import Notebook, NotebookCell
 
-TUTORIAL_NOTEBOOK_NAME = "Python Development with Cognee Tutorial 🧠"
-
-
-async def _create_tutorial_notebook(
-    user_id: UUID, session: AsyncSession, force_refresh: bool = False
-) -> None:
-    """
-    Create the default tutorial notebook for new users.
-    Dynamically fetches from: https://github.com/topoteretes/cognee/blob/notebook_tutorial/notebooks/starter_tutorial.zip
-    """
-    TUTORIAL_ZIP_URL = (
-        "https://github.com/topoteretes/cognee/raw/notebook_tutorial/notebooks/starter_tutorial.zip"
-    )
-
-    try:
-        # Create notebook from remote zip file (includes notebook + data files)
-        notebook = await Notebook.from_ipynb_zip_url(
-            zip_url=TUTORIAL_ZIP_URL,
-            owner_id=user_id,
-            notebook_filename="tutorial.ipynb",
-            name=TUTORIAL_NOTEBOOK_NAME,
-            deletable=False,
-            force=force_refresh,
-        )
-
-        # Add to session and commit
-        session.add(notebook)
-        await session.commit()
-
-    except Exception as e:
-        print(f"Failed to fetch tutorial notebook from {TUTORIAL_ZIP_URL}: {e}")
-
-        raise e
-
 
 @with_async_session
 async def create_notebook(

cognee/modules/notebooks/methods/create_tutorial_notebooks.py

@@ -0,0 +1,191 @@
+from pathlib import Path
+from uuid import NAMESPACE_OID, UUID, uuid5, uuid4
+from typing import List, Optional, Dict, Any
+import re
+import json
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from cognee.shared.logging_utils import get_logger
+from cognee.root_dir import ROOT_DIR
+
+from ..models.Notebook import Notebook, NotebookCell
+
+logger = get_logger()
+
+
+def _get_tutorials_directory() -> Path:
+    """Get the path to the tutorials directory."""
+    return ROOT_DIR / "modules" / "notebooks" / "tutorials"
+
+
+def _parse_cell_index(filename: str) -> int:
+    """Extract cell index from filename like 'cell-0.md' or 'cell-123.py'."""
+    match = re.search(r"cell-(\d+)", filename)
+    if match:
+        return int(match.group(1))
+    return -1
+
+
+def _get_cell_type(file_path: Path) -> str:
+    """Determine cell type from file extension."""
+    extension = file_path.suffix.lower()
+    if extension == ".md":
+        return "markdown"
+    elif extension == ".py":
+        return "code"
+    else:
+        raise ValueError(f"Unsupported cell file type: {extension}")
+
+
+def _extract_markdown_heading(content: str) -> str | None:
+    """Extract the first markdown heading from content."""
+    for line in content.splitlines():
+        line = line.strip()
+        # Match lines starting with one or more # followed by space and text
+        match = re.match(r"^#+\s+(.+)$", line)
+        if match:
+            return match.group(1).strip()
+    return None
+
+
+def _get_cell_name(cell_file: Path, cell_type: str, content: str) -> str:
+    """Get the appropriate name for a cell."""
+    if cell_type == "code":
+        return "Code Cell"
+    elif cell_type == "markdown":
+        heading = _extract_markdown_heading(content)
+        if heading:
+            return heading
+    # Fallback to filename stem
+    return cell_file.stem
+
+
+def _load_tutorial_cells(tutorial_dir: Path) -> List[NotebookCell]:
+    """Load all cells from a tutorial directory, sorted by cell index."""
+    cells = []
+
+    cell_files = [
+        file_path
+        for file_path in tutorial_dir.iterdir()
+        if file_path.is_file()
+        and file_path.name.startswith("cell-")
+        and file_path.suffix in [".md", ".py"]
+    ]
+
+    cell_files.sort(key=lambda file_path: _parse_cell_index(file_path.name))
+
+    for cell_file in cell_files:
+        try:
+            cell_type = _get_cell_type(cell_file)
+            content = cell_file.read_text(encoding="utf-8")
+            cell_name = _get_cell_name(cell_file, cell_type, content)
+
+            cells.append(
+                NotebookCell(
+                    id=uuid4(),
+                    type=cell_type,
+                    name=cell_name,
+                    content=content,
+                )
+            )
+        except Exception as e:
+            logger.warning(f"Failed to load cell {cell_file}: {e}")
+            continue
+
+    return cells
+
+
+def _read_tutorial_config(tutorial_dir: Path) -> Optional[Dict[str, Any]]:
+    """Read config.json from a tutorial directory if it exists."""
+    config_path = tutorial_dir / "config.json"
+    if config_path.exists():
+        try:
+            with open(config_path, "r", encoding="utf-8") as f:
+                return json.load(f)
+        except (json.JSONDecodeError, IOError) as e:
+            logger.warning(f"Failed to read config.json from {tutorial_dir}: {e}")
+            return None
+    return None
+
+
+def _format_tutorial_name(tutorial_dir_name: str) -> str:
+    """Format tutorial directory name into a readable notebook name (fallback)."""
+
+    name = tutorial_dir_name.replace("-", " ").replace("_", " ")
+    return f"{name.capitalize()} - tutorial 🧠"
+
+
+async def create_tutorial_notebooks(user_id: UUID, session: AsyncSession) -> None:
+    """
+    Create tutorial notebooks for all tutorials found in the tutorials directory.
+    Each tutorial directory will become a separate notebook.
+    """
+    try:
+        tutorials_dir = _get_tutorials_directory()
+
+        if not tutorials_dir.exists():
+            logger.warning(f"Tutorials directory not found: {tutorials_dir}")
+            return
+
+        tutorial_dirs = [
+            d for d in tutorials_dir.iterdir() if d.is_dir() and not d.name.startswith(".")
+        ]
+
+        if not tutorial_dirs:
+            logger.warning(f"No tutorial directories found in {tutorials_dir}")
+            return
+
+        notebooks_to_add = []
+
+        for tutorial_dir in tutorial_dirs:
+            try:
+                cells = _load_tutorial_cells(tutorial_dir)
+
+                if not cells:
+                    logger.warning(f"No cells found in tutorial directory: {tutorial_dir}")
+                    continue
+
+                config = _read_tutorial_config(tutorial_dir)
+
+                # Use name from config.json, or fallback to formatted directory name
+                if config and "name" in config:
+                    notebook_name = config["name"]
+                else:
+                    notebook_name = _format_tutorial_name(tutorial_dir.name)
+                    logger.warning(
+                        f"No config.json or 'name' field found in {tutorial_dir}, "
+                        f"using fallback name: {notebook_name}"
+                    )
+
+                # Use deletable flag from config.json, or default to False for tutorials
+                deletable = False
+                if config and "deletable" in config:
+                    deletable = bool(config["deletable"])
+
+                notebook_id = uuid5(NAMESPACE_OID, name=notebook_name)
+
+                notebook = Notebook(
+                    id=notebook_id,
+                    owner_id=user_id,
+                    name=notebook_name,
+                    cells=cells,
+                    deletable=deletable,
+                )
+
+                notebooks_to_add.append(notebook)
+                logger.info(f"Created tutorial notebook: {notebook_name} with {len(cells)} cells")
+
+            except Exception as e:
+                logger.error(f"Failed to create tutorial notebook from {tutorial_dir}: {e}")
+                continue
+
+        if not notebooks_to_add:
+            return
+
+        for notebook in notebooks_to_add:
+            session.add(notebook)
+
+        await session.commit()
+
+    except Exception as e:
+        logger.error(f"Failed to create tutorial notebooks for user {user_id}: {e}")

cognee/modules/notebooks/methods/get_notebooks.py

@@ -1,4 +1,4 @@
-from uuid import UUID
+from uuid import NAMESPACE_OID, UUID, uuid5
 from typing import List
 from sqlalchemy import select, and_
 from sqlalchemy.ext.asyncio import AsyncSession
@@ -6,7 +6,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
 from cognee.infrastructure.databases.relational import with_async_session
 
 from ..models.Notebook import Notebook
-from .create_notebook import _create_tutorial_notebook, TUTORIAL_NOTEBOOK_NAME
+from .create_tutorial_notebooks import create_tutorial_notebooks
 
 from cognee.shared.logging_utils import get_logger
 
@@ -19,21 +19,25 @@ async def get_notebooks(
     session: AsyncSession,
 ) -> List[Notebook]:
     # Check if tutorial notebook already exists for this user
+    tutorial_notebook_ids = [
+        uuid5(NAMESPACE_OID, name="Cognee Basics - tutorial 🧠"),
+        uuid5(NAMESPACE_OID, name="Python Development with Cognee - tutorial 🧠"),
+    ]
     tutorial_query = select(Notebook).where(
         and_(
             Notebook.owner_id == user_id,
-            Notebook.name == TUTORIAL_NOTEBOOK_NAME,
+            Notebook.id.in_(tutorial_notebook_ids),
             ~Notebook.deletable,
         )
     )
     tutorial_result = await session.execute(tutorial_query)
-    tutorial_notebook = tutorial_result.scalar_one_or_none()
+    tutorial_notebooks = tutorial_result.scalars().all()
 
-    # If tutorial notebook doesn't exist, create it
-    if tutorial_notebook is None:
-        logger.info(f"Tutorial notebook not found for user {user_id}, creating it")
+    # If tutorial notebooks don't exist, create them
+    if len(tutorial_notebooks) == 0:
+        logger.info(f"Tutorial notebooks not found for user {user_id}, creating them")
         try:
-            await _create_tutorial_notebook(user_id, session, force_refresh=False)
+            await create_tutorial_notebooks(user_id, session)
         except Exception as e:
             # Log the error but continue to return existing notebooks
             logger.error(f"Failed to create tutorial notebook for user {user_id}: {e}")

cognee/modules/notebooks/tutorials/cognee-basics/cell-1.md

@@ -0,0 +1,3 @@
+# Use Cognee to Build Your Own Knowledge Graph
+
+Cognee is a tool that allows you to build your own knowledge graph from the data you have.

cognee/modules/notebooks/tutorials/cognee-basics/cell-2.md

@@ -0,0 +1,10 @@
+# What You'll Learn in This Tutorial
+
+In this tutorial, you'll learn how to use Cognee to transform scattered data into an intelligent knowledge system that enhances your workflow.
+By the end, you'll have:
+
+- Connected disparate data sources into a unified AI memory graph
+- Built a memory layer that infers knowledge from provided data
+- Learn how to use search capabilities that combine the diverse context
+
+This tutorial demonstrates the power of knowledge graphs and retrieval-augmented generation (RAG), showing you how to build systems that learn from data and infer knowledge.

cognee/modules/notebooks/tutorials/cognee-basics/cell-3.md

@@ -0,0 +1,7 @@
+# Cognee and Its Core Operations
+
+Before we dive in, let's understand the core Cognee operations we'll be working with:
+
+- `cognee.add()` - Ingests raw data into the system
+- `cognee.cognify()` - Processes and structures data into a knowledge graph using AI
+- `cognee.search()` - Queries the knowledge graph with natural language

cognee/modules/notebooks/tutorials/cognee-basics/cell-4.py

@@ -0,0 +1,28 @@
+# Add data one by one, or pass a list to add multiple items at once
+
+await cognee.add(
+    "Harry Potter is a student at Hogwarts and belongs to Gryffindor house. \
+    He is known for defeating Voldemort and his Patronus is a stag.",
+    dataset_name="cognee-basics",
+)
+
+await cognee.add(
+    "Hermione Granger is a student at Hogwarts and also belongs to Gryffindor house. \
+    She is known for her intelligence and deep knowledge of spells. Her Patronus is an otter.",
+    dataset_name="cognee-basics",
+)
+
+await cognee.add(
+    "Severus Snape is a professor at Hogwarts who teaches Potions. \
+    He belongs to Slytherin house and was secretly loyal to Albus Dumbledore.",
+    dataset_name="cognee-basics",
+)
+
+await cognee.add(
+    [
+        "Hogwarts is a magical school located in Scotland. During Harry Potter's time at school, the headmaster was Albus Dumbledore.",
+        "A Horcrux is a dark magic object used to store a fragment of a wizard's soul. Voldemort created multiple Horcruxes to achieve immortality.",
+        "The Elder Wand is a powerful wand believed to be unbeatable. Its final known owner was Harry Potter.",
+    ],
+    dataset_name="cognee-basics",
+)