cognee 0.4.0__py3-none-any.whl → 0.5.0__py3-none-any.whl

cognee/tasks/storage/index_data_points.py

@@ -8,47 +8,58 @@ logger = get_logger("index_data_points")
 
 
 async def index_data_points(data_points: list[DataPoint]):
-    created_indexes = {}
-    index_points = {}
+    """Index data points in the vector engine by creating embeddings for specified fields.
+
+    Process:
+    1. Groups data points into a nested dict: {type_name: {field_name: [points]}}
+    2. Creates vector indexes for each (type, field) combination on first encounter
+    3. Batches points per (type, field) and creates async indexing tasks
+    4. Executes all indexing tasks in parallel for efficient embedding generation
+
+    Args:
+        data_points: List of DataPoint objects to index. Each DataPoint's metadata must
+                     contain an 'index_fields' list specifying which fields to embed.
+
+    Returns:
+        The original data_points list.
+    """
+    data_points_by_type = {}
 
     vector_engine = get_vector_engine()
 
     for data_point in data_points:
         data_point_type = type(data_point)
+        type_name = data_point_type.__name__
 
         for field_name in data_point.metadata["index_fields"]:
             if getattr(data_point, field_name, None) is None:
                 continue
 
-            index_name = f"{data_point_type.__name__}_{field_name}"
+            if type_name not in data_points_by_type:
+                data_points_by_type[type_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
-
-            if index_name not in index_points:
-                index_points[index_name] = []
+            if field_name not in data_points_by_type[type_name]:
+                await vector_engine.create_vector_index(type_name, field_name)
+                data_points_by_type[type_name][field_name] = []
 
             indexed_data_point = data_point.model_copy()
             indexed_data_point.metadata["index_fields"] = [field_name]
-            index_points[index_name].append(indexed_data_point)
+            data_points_by_type[type_name][field_name].append(indexed_data_point)
 
-    tasks: list[asyncio.Task] = []
     batch_size = vector_engine.embedding_engine.get_batch_size()
 
-    for index_name_and_field, points in index_points.items():
-        first = index_name_and_field.index("_")
-        index_name = index_name_and_field[:first]
-        field_name = index_name_and_field[first + 1 :]
+    batches = (
+        (type_name, field_name, points[i : i + batch_size])
+        for type_name, fields in data_points_by_type.items()
+        for field_name, points in fields.items()
+        for i in range(0, len(points), batch_size)
+    )
 
-        # Create embedding requests per batch to run in parallel later
-        for i in range(0, len(points), batch_size):
-            batch = points[i : i + batch_size]
-            tasks.append(
-                asyncio.create_task(vector_engine.index_data_points(index_name, field_name, batch))
-            )
+    tasks = [
+        asyncio.create_task(vector_engine.index_data_points(type_name, field_name, batch_points))
+        for type_name, field_name, batch_points in batches
+    ]
 
-    # Run all embedding requests in parallel
     await asyncio.gather(*tasks)
 
     return data_points

cognee/tasks/storage/index_graph_edges.py

@@ -1,17 +1,44 @@
-import asyncio
+from collections import Counter
+from typing import Optional, Dict, Any, List, Tuple, Union
 
 from cognee.modules.engine.utils.generate_edge_id import generate_edge_id
 from cognee.shared.logging_utils import get_logger
-from collections import Counter
-from typing import Optional, Dict, Any, List, Tuple, Union
-from cognee.infrastructure.databases.vector import get_vector_engine
 from cognee.infrastructure.databases.graph import get_graph_engine
 from cognee.modules.graph.models.EdgeType import EdgeType
 from cognee.infrastructure.databases.graph.graph_db_interface import EdgeData
+from cognee.tasks.storage.index_data_points import index_data_points
 
 logger = get_logger()
 
 
+def _get_edge_text(item: dict) -> str:
+    """Extract edge text for embedding - prefers edge_text field with fallback."""
+    if "edge_text" in item:
+        return item["edge_text"]
+
+    if "relationship_name" in item:
+        return item["relationship_name"]
+
+    return ""
+
+
+def create_edge_type_datapoints(edges_data) -> list[EdgeType]:
+    """Transform raw edge data into EdgeType datapoints."""
+    edge_texts = [
+        _get_edge_text(item)
+        for edge in edges_data
+        for item in edge
+        if isinstance(item, dict) and "relationship_name" in item
+    ]
+
+    edge_types = Counter(edge_texts)
+
+    return [
+        EdgeType(id=generate_edge_id(edge_id=text), relationship_name=text, number_of_edges=count)
+        for text, count in edge_types.items()
+    ]
+
+
 async def index_graph_edges(
     edges_data: Union[List[EdgeData], List[Tuple[str, str, str, Optional[Dict[str, Any]]]]] = None,
 ):
@@ -23,24 +50,17 @@ async def index_graph_edges(
     the `relationship_name` field.
 
     Steps:
-    1. Initialize the vector engine and graph engine.
-    2. Retrieve graph edge data and count relationship types (`relationship_name`).
-    3. Create vector indexes for `relationship_name` if they don't exist.
-    4. Transform the counted relationships into `EdgeType` objects.
-    5. Index the transformed data points in the vector engine.
+    1. Initialize the graph engine if needed and retrieve edge data.
+    2. Transform edge data into EdgeType datapoints.
+    3. Index the EdgeType datapoints using the standard indexing function.
 
     Raises:
-        RuntimeError: If initialization of the vector engine or graph engine fails.
+        RuntimeError: If initialization of the graph engine fails.
 
     Returns:
         None
     """
     try:
-        created_indexes = {}
-        index_points = {}
-
-        vector_engine = get_vector_engine()
-
         if edges_data is None:
             graph_engine = await get_graph_engine()
             _, edges_data = await graph_engine.get_graph_data()
@@ -51,47 +71,7 @@ async def index_graph_edges(
         logger.error("Failed to initialize engines: %s", e)
         raise RuntimeError("Initialization error") from e
 
-    edge_types = Counter(
-        item.get("relationship_name")
-        for edge in edges_data
-        for item in edge
-        if isinstance(item, dict) and "relationship_name" in item
-    )
-
-    for text, count in edge_types.items():
-        edge = EdgeType(
-            id=generate_edge_id(edge_id=text), relationship_name=text, number_of_edges=count
-        )
-        data_point_type = type(edge)
-
-        for field_name in edge.metadata["index_fields"]:
-            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
-
-            if index_name not in index_points:
-                index_points[index_name] = []
-
-            indexed_data_point = edge.model_copy()
-            indexed_data_point.metadata["index_fields"] = [field_name]
-            index_points[index_name].append(indexed_data_point)
-
-    # Get maximum batch size for embedding model
-    batch_size = vector_engine.embedding_engine.get_batch_size()
-    tasks: list[asyncio.Task] = []
-
-    for index_name, indexable_points in index_points.items():
-        index_name, field_name = index_name.split(".")
-
-        # Create embedding tasks to run in parallel later
-        for start in range(0, len(indexable_points), batch_size):
-            batch = indexable_points[start : start + batch_size]
-
-            tasks.append(vector_engine.index_data_points(index_name, field_name, batch))
-
-    # Start all embedding tasks and wait for completion
-    await asyncio.gather(*tasks)
+    edge_type_datapoints = create_edge_type_datapoints(edges_data)
+    await index_data_points(edge_type_datapoints)
 
     return None

cognee/tests/integration/documents/CsvDocument_test.py

@@ -0,0 +1,70 @@
+import os
+import sys
+import uuid
+import pytest
+import pathlib
+from unittest.mock import patch
+
+from cognee.modules.chunking.CsvChunker import CsvChunker
+from cognee.modules.data.processing.document_types.CsvDocument import CsvDocument
+from cognee.tests.integration.documents.AudioDocument_test import mock_get_embedding_engine
+from cognee.tests.integration.documents.async_gen_zip import async_gen_zip
+
+chunk_by_row_module = sys.modules.get("cognee.tasks.chunks.chunk_by_row")
+
+
+GROUND_TRUTH = {
+    "chunk_size_10": [
+        {"token_count": 9, "len_text": 26, "cut_type": "row_cut", "chunk_index": 0},
+        {"token_count": 6, "len_text": 29, "cut_type": "row_end", "chunk_index": 1},
+        {"token_count": 9, "len_text": 25, "cut_type": "row_cut", "chunk_index": 2},
+        {"token_count": 6, "len_text": 30, "cut_type": "row_end", "chunk_index": 3},
+    ],
+    "chunk_size_128": [
+        {"token_count": 15, "len_text": 57, "cut_type": "row_end", "chunk_index": 0},
+        {"token_count": 15, "len_text": 57, "cut_type": "row_end", "chunk_index": 1},
+    ],
+}
+
+
+@pytest.mark.parametrize(
+    "input_file,chunk_size",
+    [("example_with_header.csv", 10), ("example_with_header.csv", 128)],
+)
+@patch.object(chunk_by_row_module, "get_embedding_engine", side_effect=mock_get_embedding_engine)
+@pytest.mark.asyncio
+async def test_CsvDocument(mock_engine, input_file, chunk_size):
+    # Define file paths of test data
+    csv_file_path = os.path.join(
+        pathlib.Path(__file__).parent.parent.parent,
+        "test_data",
+        input_file,
+    )
+
+    # Define test documents
+    csv_document = CsvDocument(
+        id=uuid.uuid4(),
+        name="example_with_header.csv",
+        raw_data_location=csv_file_path,
+        external_metadata="",
+        mime_type="text/csv",
+    )
+
+    # TEST CSV
+    ground_truth_key = f"chunk_size_{chunk_size}"
+    async for ground_truth, row_data in async_gen_zip(
+        GROUND_TRUTH[ground_truth_key],
+        csv_document.read(chunker_cls=CsvChunker, max_chunk_size=chunk_size),
+    ):
+        assert ground_truth["token_count"] == row_data.chunk_size, (
+            f'{ground_truth["token_count"] = } != {row_data.chunk_size = }'
+        )
+        assert ground_truth["len_text"] == len(row_data.text), (
+            f'{ground_truth["len_text"] = } != {len(row_data.text) = }'
+        )
+        assert ground_truth["cut_type"] == row_data.cut_type, (
+            f'{ground_truth["cut_type"] = } != {row_data.cut_type = }'
+        )
+        assert ground_truth["chunk_index"] == row_data.chunk_index, (
+            f'{ground_truth["chunk_index"] = } != {row_data.chunk_index = }'
+        )

cognee/tests/integration/retrieval/test_triplet_retriever.py

@@ -0,0 +1,84 @@
+import os
+import pytest
+import pathlib
+import pytest_asyncio
+import cognee
+
+from cognee.low_level import setup
+from cognee.tasks.storage import add_data_points
+from cognee.modules.retrieval.exceptions.exceptions import NoDataError
+from cognee.modules.retrieval.triplet_retriever import TripletRetriever
+from cognee.modules.engine.models import Triplet
+
+
+@pytest_asyncio.fixture
+async def setup_test_environment_with_triplets():
+    """Set up a clean test environment with triplets."""
+    base_dir = pathlib.Path(__file__).parent.parent.parent.parent
+    system_directory_path = str(base_dir / ".cognee_system/test_triplet_retriever_context_simple")
+    data_directory_path = str(base_dir / ".data_storage/test_triplet_retriever_context_simple")
+
+    cognee.config.system_root_directory(system_directory_path)
+    cognee.config.data_root_directory(data_directory_path)
+
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+    await setup()
+
+    triplet1 = Triplet(
+        from_node_id="node1",
+        to_node_id="node2",
+        text="Alice knows Bob",
+    )
+    triplet2 = Triplet(
+        from_node_id="node2",
+        to_node_id="node3",
+        text="Bob works at Tech Corp",
+    )
+
+    triplets = [triplet1, triplet2]
+    await add_data_points(triplets)
+
+    yield
+
+    try:
+        await cognee.prune.prune_data()
+        await cognee.prune.prune_system(metadata=True)
+    except Exception:
+        pass
+
+
+@pytest_asyncio.fixture
+async def setup_test_environment_empty():
+    """Set up a clean test environment without triplets."""
+    base_dir = pathlib.Path(__file__).parent.parent.parent.parent
+    system_directory_path = str(
+        base_dir / ".cognee_system/test_triplet_retriever_context_empty_collection"
+    )
+    data_directory_path = str(
+        base_dir / ".data_storage/test_triplet_retriever_context_empty_collection"
+    )
+
+    cognee.config.system_root_directory(system_directory_path)
+    cognee.config.data_root_directory(data_directory_path)
+
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+
+    yield
+
+    try:
+        await cognee.prune.prune_data()
+        await cognee.prune.prune_system(metadata=True)
+    except Exception:
+        pass
+
+
+@pytest.mark.asyncio
+async def test_triplet_retriever_context_simple(setup_test_environment_with_triplets):
+    """Integration test: verify TripletRetriever can retrieve triplet context."""
+    retriever = TripletRetriever(top_k=5)
+
+    context = await retriever.get_context("Alice")
+
+    assert "Alice knows Bob" in context, "Failed to get Alice triplet"

cognee/tests/integration/tasks/test_add_data_points.py

@@ -0,0 +1,139 @@
+import pathlib
+import pytest
+import pytest_asyncio
+
+import cognee
+from cognee.low_level import setup
+from cognee.infrastructure.engine import DataPoint
+from cognee.tasks.storage.add_data_points import add_data_points
+from cognee.tasks.storage.exceptions import InvalidDataPointsInAddDataPointsError
+from cognee.infrastructure.databases.graph import get_graph_engine
+
+
+class Person(DataPoint):
+    name: str
+    age: int
+    metadata: dict = {"index_fields": ["name"]}
+
+
+class Company(DataPoint):
+    name: str
+    industry: str
+    metadata: dict = {"index_fields": ["name", "industry"]}
+
+
+@pytest_asyncio.fixture
+async def clean_test_environment():
+    """Set up a clean test environment for add_data_points tests."""
+    base_dir = pathlib.Path(__file__).parent.parent.parent.parent
+    system_directory_path = str(base_dir / ".cognee_system/test_add_data_points_integration")
+    data_directory_path = str(base_dir / ".data_storage/test_add_data_points_integration")
+
+    cognee.config.system_root_directory(system_directory_path)
+    cognee.config.data_root_directory(data_directory_path)
+
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+    await setup()
+
+    yield
+
+    try:
+        await cognee.prune.prune_data()
+        await cognee.prune.prune_system(metadata=True)
+    except Exception:
+        pass
+
+
+@pytest.mark.asyncio
+async def test_add_data_points_comprehensive(clean_test_environment):
+    """Comprehensive integration test for add_data_points functionality."""
+
+    person1 = Person(name="Alice", age=30)
+    person2 = Person(name="Bob", age=25)
+    result = await add_data_points([person1, person2])
+
+    assert result == [person1, person2]
+    assert len(result) == 2
+
+    graph_engine = await get_graph_engine()
+    nodes, edges = await graph_engine.get_graph_data()
+    assert len(nodes) >= 2
+
+    result_empty = await add_data_points([])
+    assert result_empty == []
+
+    person3 = Person(name="Charlie", age=35)
+    person4 = Person(name="Diana", age=32)
+    custom_edge = (str(person3.id), str(person4.id), "knows", {"edge_text": "friends with"})
+
+    result_custom = await add_data_points([person3, person4], custom_edges=[custom_edge])
+    assert len(result_custom) == 2
+
+    nodes, edges = await graph_engine.get_graph_data()
+    assert len(edges) == 1
+    assert len(nodes) == 4
+
+    class Employee(DataPoint):
+        name: str
+        works_at: Company
+        metadata: dict = {"index_fields": ["name"]}
+
+    company = Company(name="TechCorp", industry="Technology")
+    employee = Employee(name="Eve", works_at=company)
+
+    result_rel = await add_data_points([employee])
+    assert len(result_rel) == 1
+
+    nodes, edges = await graph_engine.get_graph_data()
+    assert len(nodes) == 6
+    assert len(edges) == 2
+
+    person5 = Person(name="Frank", age=40)
+    person6 = Person(name="Grace", age=38)
+    triplet_edge = (str(person5.id), str(person6.id), "married_to", {"edge_text": "is married to"})
+
+    result_triplet = await add_data_points(
+        [person5, person6], custom_edges=[triplet_edge], embed_triplets=True
+    )
+    assert len(result_triplet) == 2
+
+    nodes, edges = await graph_engine.get_graph_data()
+    assert len(nodes) == 8
+    assert len(edges) == 3
+
+    batch1 = [Person(name="Leo", age=25), Person(name="Mia", age=30)]
+    batch2 = [Person(name="Noah", age=35), Person(name="Olivia", age=40)]
+
+    result_batch1 = await add_data_points(batch1)
+    result_batch2 = await add_data_points(batch2)
+
+    assert len(result_batch1) == 2
+    assert len(result_batch2) == 2
+
+    nodes, edges = await graph_engine.get_graph_data()
+    assert len(nodes) == 12
+    assert len(edges) == 3
+
+    person7 = Person(name="Paul", age=33)
+    person8 = Person(name="Quinn", age=31)
+    edge1 = (str(person7.id), str(person8.id), "colleague_of", {"edge_text": "works with"})
+    edge2 = (str(person8.id), str(person7.id), "colleague_of", {"edge_text": "works with"})
+
+    result_bi = await add_data_points([person7, person8], custom_edges=[edge1, edge2])
+    assert len(result_bi) == 2
+
+    nodes, edges = await graph_engine.get_graph_data()
+    assert len(nodes) == 14
+    assert len(edges) == 5
+
+    person_invalid = Person(name="Invalid", age=50)
+    with pytest.raises(InvalidDataPointsInAddDataPointsError, match="must be a list"):
+        await add_data_points(person_invalid)
+
+    with pytest.raises(InvalidDataPointsInAddDataPointsError, match="must be a DataPoint"):
+        await add_data_points(["not", "datapoints"])
+
+    final_nodes, final_edges = await graph_engine.get_graph_data()
+    assert len(final_nodes) == 14
+    assert len(final_edges) == 5

cognee/tests/integration/tasks/test_get_triplet_datapoints.py

@@ -0,0 +1,69 @@
+import os
+import pathlib
+import pytest
+import pytest_asyncio
+from unittest.mock import AsyncMock, patch
+
+import cognee
+from cognee.tasks.memify.get_triplet_datapoints import get_triplet_datapoints
+from cognee.modules.engine.models import Triplet
+
+
+@pytest_asyncio.fixture
+async def setup_test_environment():
+    """Set up a clean test environment with a simple graph."""
+    base_dir = pathlib.Path(__file__).parent.parent.parent.parent
+    data_directory_path = str(base_dir / ".data_storage/test_get_triplet_datapoints_integration")
+    cognee_directory_path = str(base_dir / ".cognee_system/test_get_triplet_datapoints_integration")
+
+    cognee.config.data_root_directory(data_directory_path)
+    cognee.config.system_root_directory(cognee_directory_path)
+
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+
+    dataset_name = "test_triplets"
+
+    text = "Volkswagen is a german car manufacturer from Wolfsburg. They produce different models such as Golf, Polo and Touareg."
+    await cognee.add(text, dataset_name)
+    await cognee.cognify([dataset_name])
+
+    yield dataset_name
+
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+
+
+@pytest.mark.asyncio
+async def test_get_triplet_datapoints_integration(setup_test_environment):
+    """Integration test: verify get_triplet_datapoints works with real graph data."""
+
+    from cognee.infrastructure.databases.graph import get_graph_engine
+
+    graph_engine = await get_graph_engine()
+
+    if not hasattr(graph_engine, "get_triplets_batch"):
+        pytest.skip("Graph engine does not support get_triplets_batch")
+
+    triplets = []
+    with patch(
+        "cognee.tasks.memify.get_triplet_datapoints.index_data_points", new_callable=AsyncMock
+    ):
+        async for triplet in get_triplet_datapoints([{}], triplets_batch_size=10):
+            triplets.append(triplet)
+
+    nodes, edges = await graph_engine.get_graph_data()
+
+    if len(edges) > 0 and len(triplets) == 0:
+        test_triplets = await graph_engine.get_triplets_batch(offset=0, limit=10)
+        if len(test_triplets) == 0:
+            pytest.fail(
+                f"Edges exist in graph ({len(edges)} edges) but get_triplets_batch found none. "
+                f"This indicates the query pattern may not match the graph structure."
+            )
+
+    for triplet in triplets:
+        assert isinstance(triplet, Triplet), "Each item should be a Triplet instance"
+        assert triplet.from_node_id, "Triplet should have from_node_id"
+        assert triplet.to_node_id, "Triplet should have to_node_id"
+        assert triplet.text, "Triplet should have embeddable text"

cognee/tests/integration/web_url_crawler/test_default_url_crawler.py

@@ -5,7 +5,7 @@ from cognee.tasks.web_scraper import DefaultUrlCrawler
 @pytest.mark.asyncio
 async def test_fetch():
     crawler = DefaultUrlCrawler()
-    url = "https://en.wikipedia.org/wiki/Large_language_model"
+    url = "http://example.com/"
     results = await crawler.fetch_urls(url)
     assert len(results) == 1
     assert isinstance(results, dict)

cognee/tests/integration/web_url_crawler/test_tavily_crawler.py

@@ -11,7 +11,7 @@ skip_in_ci = pytest.mark.skipif(
 @skip_in_ci
 @pytest.mark.asyncio
 async def test_fetch():
-    url = "https://en.wikipedia.org/wiki/Large_language_model"
+    url = "http://example.com/"
     results = await fetch_with_tavily(url)
     assert isinstance(results, dict)
     assert len(results) == 1