cognee 0.2.4__py3-none-any.whl → 0.3.0.dev0__py3-none-any.whl

cognee/tests/test_temporal_graph.py

@@ -0,0 +1,167 @@
+import asyncio
+import cognee
+from cognee.modules.retrieval.temporal_retriever import TemporalRetriever
+
+from cognee.shared.logging_utils import setup_logging, INFO
+from cognee.tasks.temporal_graph.models import Timestamp
+from cognee.api.v1.search import SearchType
+from cognee.shared.logging_utils import get_logger
+from cognee.infrastructure.databases.graph.get_graph_engine import get_graph_engine
+from collections import Counter
+from cognee.modules.engine.utils.generate_timestamp_datapoint import date_to_int
+
+logger = get_logger()
+
+biography_1 = """
+    Attaphol Buspakom Attaphol Buspakom ( ; ) , nicknamed Tak ( ; ) ; 1 October 1962 – 16 April 2015 ) was a Thai national and football coach . He was given the role at Muangthong United and Buriram United after TTM Samut Sakhon folded after the 2009 season . He played for the Thailand national football team , appearing in several FIFA World Cup qualifying matches .
+
+    Club career .
+    Attaphol began his career as a player at Thai Port FC Authority of Thailand in 1985 . In his first year , he won his first championship with the club . He played for the club until 1989 and in 1987 also won the Queens Cup . He then moved to Malaysia for two seasons for Pahang FA , then return to Thailand to his former club . His time from 1991 to 1994 was marked by less success than in his first stay at Port Authority . From 1994 to 1996 he played for Pahang again and this time he was able to win with the club , the Malaysia Super League and also reached the final of the Malaysia Cup and the Malaysia FA Cup . Both cup finals but lost . Back in Thailand , he let end his playing career at FC Stock Exchange of Thailand , with which he once again runner‑up in 1996-97 . In 1998 , he finished his career .
+
+    International career .
+    For the Thailand national football team Attaphol played between 1985 and 1998 a total of 85 games and scored 13 results . In 1992 , he participated with the team in the finals of the Asian Cup . He also stood in various cadres to qualifications to FIFA World Cup .
+
+    Coaching career .
+    Bec Tero Sasana .
+    In BEC Tero Sasana F.C . began his coaching career in 2001 for him , first as assistant coach . He took over the reigning champions of the Thai League T1 , after his predecessor Pichai Pituwong resigned from his post . It was his first coach station and he had the difficult task of leading the club through the new AFC Champions League . He could accomplish this task with flying colors and even led the club to the finals . The finale , then still played in home and away matches , was lost with 1:2 at the end against Al Ain FC . Attaphol is and was next to Charnwit Polcheewin the only coach who managed a club from Thailand to lead to the final of the AFC Champions League . 2002-03 and 2003-04 he won with the club also two runner‑up . In his team , which reached the final of the Champions League , were a number of exceptional players like Therdsak Chaiman , Worrawoot Srimaka , Dusit Chalermsan and Anurak Srikerd .
+
+    Geylang United / Krung Thai Bank .
+    In 2006 , he went to Singapore in the S‑League to Geylang United He was released after a few months due to lack of success . In 2008 , he took over as coach at Krung Thai Bank F.C. , where he had almost a similar task , as a few years earlier by BEC‑Tero . As vice‑champion of the club was also qualified for the AFC Champions League . However , he failed to lead the team through the group stage of the season 2008 and beyond . With the Kashima Antlers of Japan and Beijing Guoan F.C . athletic competition was too great . One of the highlights was put under his leadership , yet the club . In the group match against the Vietnam club Nam Dinh F.C . his team won with 9-1 , but also lost four weeks later with 1-8 against Kashima Antlers . At the end of the National Football League season , he reached the Krung Thai 6th Table space . The Erstligalizenz the club was sold at the end of the season at the Bangkok Glass F.C. . Attaphol finished his coaching career with the club and accepted an offer of TTM Samutsakorn . After only a short time in office
+
+    Muangthong United .
+    In 2009 , he received an offer from Muangthong United F.C. , which he accepted and changed . He can champion Muang Thong United for 2009 Thai Premier League and Attaphol won Coach of The year for Thai Premier League and he was able to lead Muang Thong United to play AFC Champions League qualifying play‑off for the first in the clubs history .
+
+    Buriram United .
+    In 2010 Buspakom moved from Muangthong United to Buriram United F.C. . He received Coach of the Month in Thai Premier League 2 time in June and October . In 2011 , he led Buriram United win 2011 Thai Premier League second time for club and set a record with the most points in the Thai League T1 for 85 point and He led Buriram win 2011 Thai FA Cup by beat Muangthong United F.C . 1‑0 and he led Buriram win 2011 Thai League Cup by beat Thai Port F.C . 2‑0 . In 2012 , he led Buriram United to the 2012 AFC Champions League group stage . Buriram along with Guangzhou Evergrande F.C . from China , Kashiwa Reysol from Japan and Jeonbuk Hyundai Motors which are all champions from their country . In the first match of Buriram they beat Kashiwa 3‑2 and Second Match they beat Guangzhou 1‑2 at the Tianhe Stadium . Before losing to Jeonbuk 0‑2 and 3‑2 with lose Kashiwa and Guangzhou 1‑0 and 1‑2 respectively and Thai Premier League Attaphol lead Buriram end 4th for table with win 2012 Thai FA Cup and 2012 Thai League Cup .
+
+    Bangkok Glass .
+    In 2013 , he moved from Buriram United to Bangkok Glass F.C. .
+
+    Individual
+    - Thai Premier League Coach of the Year ( 3 ) : 2001-02 , 2009 , 2013
+    """
+
+
+biography_2 = """
+    Arnulf Øverland Ole Peter Arnulf Øverland ( 27 April 1889 – 25 March 1968 ) was a Norwegian poet and artist . He is principally known for his poetry which served to inspire the Norwegian resistance movement during the German occupation of Norway during World War II .
+
+    Biography .
+    Øverland was born in Kristiansund and raised in Bergen . His parents were Peter Anton Øverland ( 1852–1906 ) and Hanna Hage ( 1854–1939 ) . The early death of his father , left the family economically stressed . He was able to attend Bergen Cathedral School and in 1904 Kristiania Cathedral School . He graduated in 1907 and for a time studied philology at University of Kristiania . Øverland published his first collection of poems ( 1911 ) .
+
+    Øverland became a communist sympathizer from the early 1920s and became a member of Mot Dag . He also served as chairman of the Norwegian Students Society 1923–28 . He changed his stand in 1937 , partly as an expression of dissent against the ongoing Moscow Trials . He was an avid opponent of Nazism and in 1936 he wrote the poem Du må ikke sove which was printed in the journal Samtiden . It ends with . ( I thought: : Something is imminent . Our era is over – Europe’s on fire! ) . Probably the most famous line of the poem is ( You mustnt endure so well the injustice that doesnt affect you yourself! )
+
+    During the German occupation of Norway from 1940 in World War II , he wrote to inspire the Norwegian resistance movement . He wrote a series of poems which were clandestinely distributed , leading to the arrest of both him and his future wife Margrete Aamot Øverland in 1941 . Arnulf Øverland was held first in the prison camp of Grini before being transferred to Sachsenhausen concentration camp in Germany . He spent a four‑year imprisonment until the liberation of Norway in 1945 . His poems were later collected in Vi overlever alt and published in 1945 .
+
+    Øverland played an important role in the Norwegian language struggle in the post‑war era . He became a noted supporter for the conservative written form of Norwegian called Riksmål , he was president of Riksmålsforbundet ( an organization in support of Riksmål ) from 1947 to 1956 . In addition , Øverland adhered to the traditionalist style of writing , criticising modernist poetry on several occasions . His speech Tungetale fra parnasset , published in Arbeiderbladet in 1954 , initiated the so‑called Glossolalia debate .
+
+    Personal life .
+    In 1918 he had married the singer Hildur Arntzen ( 1888–1957 ) . Their marriage was dissolved in 1939 . In 1940 , he married Bartholine Eufemia Leganger ( 1903–1995 ) . They separated shortly after , and were officially divorced in 1945 . Øverland was married to journalist Margrete Aamot Øverland ( 1913–1978 ) during June 1945 . In 1946 , the Norwegian Parliament arranged for Arnulf and Margrete Aamot Øverland to reside at the Grotten . He lived there until his death in 1968 and she lived there for another ten years until her death in 1978 . Arnulf Øverland was buried at Vår Frelsers Gravlund in Oslo . Joseph Grimeland designed the bust of Arnulf Øverland ( bronze , 1970 ) at his grave site .
+
+    Selected Works .
+    - Den ensomme fest ( 1911 )
+    - Berget det blå ( 1927 )
+    - En Hustavle ( 1929 )
+    - Den røde front ( 1937 )
+    - Vi overlever alt ( 1945 )
+    - Sverdet bak døren ( 1956 )
+    - Livets minutter ( 1965 )
+
+    Awards .
+    - Gyldendals Endowment ( 1935 )
+    - Dobloug Prize ( 1951 )
+    - Mads Wiel Nygaards legat ( 1961 )
+    """
+
+
+async def main():
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+
+    await cognee.add([biography_1, biography_2])
+
+    await cognee.cognify(temporal_cognify=True)
+
+    graph_engine = await get_graph_engine()
+    graph = await graph_engine.get_graph_data()
+
+    type_counts = Counter(node_data[1].get("type", {}) for node_data in graph[0])
+
+    edge_type_counts = Counter(edge_type[2] for edge_type in graph[1])
+
+    # Graph structure test
+    assert type_counts.get("TextDocument", 0) == 2, (
+        f"Expected exactly one TextDocument, but found {type_counts.get('TextDocument', 0)}"
+    )
+
+    assert type_counts.get("DocumentChunk", 0) == 2, (
+        f"Expected exactly one DocumentChunk, but found {type_counts.get('DocumentChunk', 0)}"
+    )
+
+    assert type_counts.get("Entity", 0) >= 20, (
+        f"Expected multiple entities (assert is set to 20), but found {type_counts.get('Entity', 0)}"
+    )
+
+    assert type_counts.get("EntityType", 0) >= 2, (
+        f"Expected multiple entity types, but found {type_counts.get('EntityType', 0)}"
+    )
+
+    assert type_counts.get("Event", 0) >= 20, (
+        f"Expected multiple events (assert is set to 20), but found {type_counts.get('Event', 0)}"
+    )
+
+    assert type_counts.get("Timestamp", 0) >= 20, (
+        f"Expected multiple timestamps (assert is set to 20), but found {type_counts.get('Timestamp', 0)}"
+    )
+
+    assert type_counts.get("Interval", 0) >= 2, (
+        f"Expected multiple intervals, but found {type_counts.get('Interval', 0)}"
+    )
+
+    assert edge_type_counts.get("contains", 0) >= 20, (
+        f"Expected multiple 'contains' edge, but found {edge_type_counts.get('contains', 0)}"
+    )
+
+    assert edge_type_counts.get("is_a", 0) >= 20, (
+        f"Expected multiple 'is_a' edge, but found {edge_type_counts.get('is_a', 0)}"
+    )
+
+    assert edge_type_counts.get("during", 0) == type_counts.get("Interval", 0), (
+        "Expected the same amount of during and interval objects in the graph"
+    )
+
+    assert edge_type_counts.get("during", 0) == type_counts.get("Interval", 0), (
+        "Expected the same amount of during and interval objects in the graph"
+    )
+
+    assert edge_type_counts.get("time_from", 0) == type_counts.get("Interval", 0), (
+        "Expected the same amount of time_from and interval objects in the graph"
+    )
+
+    assert edge_type_counts.get("time_to", 0) == type_counts.get("Interval", 0), (
+        "Expected the same amount of time_to and interval objects in the graph"
+    )
+
+    retriever = TemporalRetriever()
+
+    result_before = await retriever.extract_time_from_query("What happened before 1890?")
+
+    assert result_before[0] is None
+
+    result_after = await retriever.extract_time_from_query("What happened after 1891?")
+
+    assert result_after[1] is None
+
+    result_between = await retriever.extract_time_from_query("What happened between 1890 and 1900?")
+
+    assert result_between[1]
+    assert result_between[0]
+
+
+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())

cognee/tests/unit/api/__init__.py

@@ -0,0 +1 @@
+# Test package for API tests

cognee/tests/unit/api/test_conditional_authentication_endpoints.py

@@ -0,0 +1,246 @@
+import pytest
+from unittest.mock import patch, AsyncMock, MagicMock
+from uuid import uuid4
+from fastapi.testclient import TestClient
+from types import SimpleNamespace
+import importlib
+
+from cognee.api.client import app
+
+
+# Fixtures for reuse across test classes
+@pytest.fixture
+def mock_default_user():
+    """Mock default user for testing."""
+    return SimpleNamespace(
+        id=uuid4(), email="default@example.com", is_active=True, tenant_id=uuid4()
+    )
+
+
+@pytest.fixture
+def mock_authenticated_user():
+    """Mock authenticated user for testing."""
+    from cognee.modules.users.models import User
+
+    return User(
+        id=uuid4(),
+        email="auth@example.com",
+        hashed_password="hashed",
+        is_active=True,
+        is_verified=True,
+        tenant_id=uuid4(),
+    )
+
+
+gau_mod = importlib.import_module("cognee.modules.users.methods.get_authenticated_user")
+
+
+class TestConditionalAuthenticationEndpoints:
+    """Test that API endpoints work correctly with conditional authentication."""
+
+    @pytest.fixture
+    def client(self):
+        """Create a test client."""
+        return TestClient(app)
+
+    def test_health_endpoint_no_auth_required(self, client):
+        """Test that health endpoint works without authentication."""
+        response = client.get("/health")
+        assert response.status_code in [200, 503]  # 503 is also acceptable for health checks
+
+    def test_root_endpoint_no_auth_required(self, client):
+        """Test that root endpoint works without authentication."""
+        response = client.get("/")
+        assert response.status_code == 200
+        assert response.json() == {"message": "Hello, World, I am alive!"}
+
+    @patch(
+        "cognee.api.client.REQUIRE_AUTHENTICATION",
+        False,
+    )
+    def test_openapi_schema_no_global_security(self, client):
+        """Test that OpenAPI schema doesn't require global authentication."""
+        response = client.get("/openapi.json")
+        assert response.status_code == 200
+
+        schema = response.json()
+
+        # Should not have global security requirement
+        global_security = schema.get("security", [])
+        assert global_security == []
+
+        # But should still have security schemes defined
+        security_schemes = schema.get("components", {}).get("securitySchemes", {})
+        assert "BearerAuth" in security_schemes
+        assert "CookieAuth" in security_schemes
+
+    @patch("cognee.api.v1.add.add")
+    @patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
+    @patch(
+        "cognee.api.client.REQUIRE_AUTHENTICATION",
+        False,
+    )
+    def test_add_endpoint_with_conditional_auth(
+        self, mock_get_default_user, mock_add, client, mock_default_user
+    ):
+        """Test add endpoint works with conditional authentication."""
+        mock_get_default_user.return_value = mock_default_user
+        mock_add.return_value = MagicMock(
+            model_dump=lambda: {"status": "success", "pipeline_run_id": str(uuid4())}
+        )
+
+        # Test file upload without authentication
+        files = {"data": ("test.txt", b"test content", "text/plain")}
+        form_data = {"datasetName": "test_dataset"}
+
+        response = client.post("/api/v1/add", files=files, data=form_data)
+
+        assert mock_get_default_user.call_count == 1
+
+        # Core test: authentication is not required (should not get 401)
+        assert response.status_code != 401
+
+    @patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
+    @patch(
+        "cognee.api.client.REQUIRE_AUTHENTICATION",
+        False,
+    )
+    def test_conditional_authentication_works_with_current_environment(
+        self, mock_get_default_user, client
+    ):
+        """Test that conditional authentication works with the current environment setup."""
+        # Since REQUIRE_AUTHENTICATION defaults to "false", we expect endpoints to work without auth
+        # This tests the actual integration behavior
+
+        mock_get_default_user.return_value = SimpleNamespace(
+            id=uuid4(), email="default@example.com", is_active=True, tenant_id=uuid4()
+        )
+
+        files = {"data": ("test.txt", b"test content", "text/plain")}
+        form_data = {"datasetName": "test_dataset"}
+
+        response = client.post("/api/v1/add", files=files, data=form_data)
+
+        assert mock_get_default_user.call_count == 1
+
+        # Core test: authentication is not required (should not get 401)
+        assert response.status_code != 401
+        # Note: This test verifies conditional authentication works in the current environment
+
+
+class TestConditionalAuthenticationBehavior:
+    """Test the behavior of conditional authentication across different endpoints."""
+
+    @pytest.fixture
+    def client(self):
+        return TestClient(app)
+
+    @pytest.mark.parametrize(
+        "endpoint,method",
+        [
+            ("/api/v1/search", "GET"),
+            ("/api/v1/datasets", "GET"),
+        ],
+    )
+    @patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
+    def test_get_endpoints_work_without_auth(
+        self, mock_get_default, client, endpoint, method, mock_default_user
+    ):
+        """Test that GET endpoints work without authentication (with current environment)."""
+        mock_get_default.return_value = mock_default_user
+
+        if method == "GET":
+            response = client.get(endpoint)
+        elif method == "POST":
+            response = client.post(endpoint, json={})
+
+        assert mock_get_default.call_count == 1
+
+        # Should not return 401 Unauthorized (authentication is optional by default)
+        assert response.status_code != 401
+
+        # May return other errors due to missing data/config, but not auth errors
+        if response.status_code >= 400:
+            # Check that it's not an authentication error
+            try:
+                error_detail = response.json().get("detail", "")
+                assert "authenticate" not in error_detail.lower()
+                assert "unauthorized" not in error_detail.lower()
+            except Exception:
+                pass  # If response is not JSON, that's fine
+
+    gsm_mod = importlib.import_module("cognee.modules.settings.get_settings")
+
+    @patch.object(gsm_mod, "get_vectordb_config")
+    @patch.object(gsm_mod, "get_llm_config")
+    @patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
+    def test_settings_endpoint_integration(
+        self, mock_get_default, mock_llm_config, mock_vector_config, client, mock_default_user
+    ):
+        """Test that settings endpoint integration works with conditional authentication."""
+        mock_get_default.return_value = mock_default_user
+
+        # Mock configurations to avoid validation errors
+        mock_llm_config.return_value = SimpleNamespace(
+            llm_provider="openai",
+            llm_model="gpt-4o",
+            llm_endpoint=None,
+            llm_api_version=None,
+            llm_api_key="test_key_1234567890",
+        )
+
+        mock_vector_config.return_value = SimpleNamespace(
+            vector_db_provider="lancedb",
+            vector_db_url="localhost:5432",  # Must be string, not None
+            vector_db_key="test_vector_key",
+        )
+
+        response = client.get("/api/v1/settings")
+
+        assert mock_get_default.call_count == 1
+
+        # Core test: authentication is not required (should not get 401)
+        assert response.status_code != 401
+        # Note: This test verifies conditional authentication works for settings endpoint
+
+
+class TestConditionalAuthenticationErrorHandling:
+    """Test error handling in conditional authentication."""
+
+    @pytest.fixture
+    def client(self):
+        return TestClient(app)
+
+    @patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
+    def test_get_default_user_fails(self, mock_get_default, client):
+        """Test behavior when get_default_user fails (with current environment)."""
+        mock_get_default.side_effect = Exception("Database connection failed")
+
+        # The error should propagate - either as a 500 error or as an exception
+        files = {"data": ("test.txt", b"test content", "text/plain")}
+        form_data = {"datasetName": "test_dataset"}
+
+        # Test that the exception is properly converted to HTTP 500
+        response = client.post("/api/v1/add", files=files, data=form_data)
+
+        # Should return HTTP 500 Internal Server Error when get_default_user fails
+        assert response.status_code == 500
+
+        # Check that the error message is informative
+        error_detail = response.json().get("detail", "")
+        assert "Failed to create default user" in error_detail
+        # The exact error message may vary depending on the actual database connection
+        # The important thing is that we get a 500 error when user creation fails
+
+    def test_current_environment_configuration(self):
+        """Test that current environment configuration is working properly."""
+        # This tests the actual module state without trying to change it
+        from cognee.modules.users.methods.get_authenticated_user import (
+            REQUIRE_AUTHENTICATION,
+        )
+
+        # Should be a boolean value (the parsing logic works)
+        assert isinstance(REQUIRE_AUTHENTICATION, bool)
+
+        # In default environment, should be False
+        assert not REQUIRE_AUTHENTICATION

cognee/tests/unit/modules/retrieval/chunks_retriever_test.py

@@ -1,7 +1,7 @@
 import os
 import pytest
 import pathlib
-
+from typing import List
 import cognee
 from cognee.low_level import setup
 from cognee.tasks.storage import add_data_points
@@ -10,6 +10,20 @@ from cognee.modules.chunking.models import DocumentChunk
 from cognee.modules.data.processing.document_types import TextDocument
 from cognee.modules.retrieval.exceptions.exceptions import NoDataError
 from cognee.modules.retrieval.chunks_retriever import ChunksRetriever
+from cognee.infrastructure.engine import DataPoint
+from cognee.modules.data.processing.document_types import Document
+from cognee.modules.engine.models import Entity
+
+
+class DocumentChunkWithEntities(DataPoint):
+    text: str
+    chunk_size: int
+    chunk_index: int
+    cut_type: str
+    is_part_of: Document
+    contains: List[Entity] = None
+
+    metadata: dict = {"index_fields": ["text"]}
 
 
 class TestChunksRetriever:
@@ -179,7 +193,9 @@ class TestChunksRetriever:
             await retriever.get_context("Christina Mayer")
 
         vector_engine = get_vector_engine()
-        await vector_engine.create_collection("DocumentChunk_text", payload_schema=DocumentChunk)
+        await vector_engine.create_collection(
+            "DocumentChunk_text", payload_schema=DocumentChunkWithEntities
+        )
 
         context = await retriever.get_context("Christina Mayer")
         assert len(context) == 0, "Found chunks when none should exist"

cognee/tests/unit/modules/retrieval/graph_completion_retriever_context_extension_test.py

@@ -6,6 +6,7 @@ from typing import Optional, Union
 import cognee
 from cognee.low_level import setup, DataPoint
 from cognee.tasks.storage import add_data_points
+from cognee.modules.graph.utils import resolve_edges_to_text
 from cognee.infrastructure.databases.exceptions import DatabaseNotCreatedError
 from cognee.modules.retrieval.graph_completion_context_extension_retriever import (
     GraphCompletionContextExtensionRetriever,
@@ -51,17 +52,15 @@ class TestGraphCompletionWithContextExtensionRetriever:
 
         retriever = GraphCompletionContextExtensionRetriever()
 
-        context, _ = await retriever.get_context("Who works at Canva?")
+        context = await resolve_edges_to_text(await retriever.get_context("Who works at Canva?"))
 
         assert "Mike Broski --[works_for]--> Canva" in context, "Failed to get Mike Broski"
         assert "Christina Mayer --[works_for]--> Canva" in context, "Failed to get Christina Mayer"
 
         answer = await retriever.get_completion("Who works at Canva?")
 
-        assert isinstance(answer, list), f"Expected list, got {type(answer).__name__}"
-        assert all(isinstance(item, str) and item.strip() for item in answer), (
-            "Answer must contain only non-empty strings"
-        )
+        assert isinstance(answer, str), f"Expected string, got {type(answer).__name__}"
+        assert answer.strip(), "Answer must contain only non-empty strings"
 
     @pytest.mark.asyncio
     async def test_graph_completion_extension_context_complex(self):
@@ -129,7 +128,9 @@ class TestGraphCompletionWithContextExtensionRetriever:
 
         retriever = GraphCompletionContextExtensionRetriever(top_k=20)
 
-        context, _ = await retriever.get_context("Who works at Figma?")
+        context = await resolve_edges_to_text(
+            await retriever.get_context("Who works at Figma and drives Tesla?")
+        )
 
         print(context)
 
@@ -139,10 +140,8 @@ class TestGraphCompletionWithContextExtensionRetriever:
 
         answer = await retriever.get_completion("Who works at Figma?")
 
-        assert isinstance(answer, list), f"Expected list, got {type(answer).__name__}"
-        assert all(isinstance(item, str) and item.strip() for item in answer), (
-            "Answer must contain only non-empty strings"
-        )
+        assert isinstance(answer, str), f"Expected string, got {type(answer).__name__}"
+        assert answer.strip(), "Answer must contain only non-empty strings"
 
     @pytest.mark.asyncio
     async def test_get_graph_completion_extension_context_on_empty_graph(self):
@@ -167,12 +166,10 @@ class TestGraphCompletionWithContextExtensionRetriever:
 
         await setup()
 
-        context, _ = await retriever.get_context("Who works at Figma?")
-        assert context == "", "Context should be empty on an empty graph"
+        context = await retriever.get_context("Who works at Figma?")
+        assert context == [], "Context should be empty on an empty graph"
 
         answer = await retriever.get_completion("Who works at Figma?")
 
-        assert isinstance(answer, list), f"Expected list, got {type(answer).__name__}"
-        assert all(isinstance(item, str) and item.strip() for item in answer), (
-            "Answer must contain only non-empty strings"
-        )
+        assert isinstance(answer, str), f"Expected string, got {type(answer).__name__}"
+        assert answer.strip(), "Answer must contain only non-empty strings"

cognee/tests/unit/modules/retrieval/graph_completion_retriever_cot_test.py

@@ -5,6 +5,7 @@ from typing import Optional, Union
 
 import cognee
 from cognee.low_level import setup, DataPoint
+from cognee.modules.graph.utils import resolve_edges_to_text
 from cognee.tasks.storage import add_data_points
 from cognee.infrastructure.databases.exceptions import DatabaseNotCreatedError
 from cognee.modules.retrieval.graph_completion_cot_retriever import GraphCompletionCotRetriever
@@ -47,17 +48,15 @@ class TestGraphCompletionCoTRetriever:
 
         retriever = GraphCompletionCotRetriever()
 
-        context, _ = await retriever.get_context("Who works at Canva?")
+        context = await resolve_edges_to_text(await retriever.get_context("Who works at Canva?"))
 
         assert "Mike Broski --[works_for]--> Canva" in context, "Failed to get Mike Broski"
         assert "Christina Mayer --[works_for]--> Canva" in context, "Failed to get Christina Mayer"
 
         answer = await retriever.get_completion("Who works at Canva?")
 
-        assert isinstance(answer, list), f"Expected list, got {type(answer).__name__}"
-        assert all(isinstance(item, str) and item.strip() for item in answer), (
-            "Answer must contain only non-empty strings"
-        )
+        assert isinstance(answer, str), f"Expected string, got {type(answer).__name__}"
+        assert answer.strip(), "Answer must contain only non-empty strings"
 
     @pytest.mark.asyncio
     async def test_graph_completion_cot_context_complex(self):
@@ -124,7 +123,7 @@ class TestGraphCompletionCoTRetriever:
 
         retriever = GraphCompletionCotRetriever(top_k=20)
 
-        context, _ = await retriever.get_context("Who works at Figma?")
+        context = await resolve_edges_to_text(await retriever.get_context("Who works at Figma?"))
 
         print(context)
 
@@ -134,10 +133,8 @@ class TestGraphCompletionCoTRetriever:
 
         answer = await retriever.get_completion("Who works at Figma?")
 
-        assert isinstance(answer, list), f"Expected list, got {type(answer).__name__}"
-        assert all(isinstance(item, str) and item.strip() for item in answer), (
-            "Answer must contain only non-empty strings"
-        )
+        assert isinstance(answer, str), f"Expected string, got {type(answer).__name__}"
+        assert answer.strip(), "Answer must contain only non-empty strings"
 
     @pytest.mark.asyncio
     async def test_get_graph_completion_cot_context_on_empty_graph(self):
@@ -162,12 +159,10 @@ class TestGraphCompletionCoTRetriever:
 
         await setup()
 
-        context, _ = await retriever.get_context("Who works at Figma?")
-        assert context == "", "Context should be empty on an empty graph"
+        context = await retriever.get_context("Who works at Figma?")
+        assert context == [], "Context should be empty on an empty graph"
 
         answer = await retriever.get_completion("Who works at Figma?")
 
-        assert isinstance(answer, list), f"Expected list, got {type(answer).__name__}"
-        assert all(isinstance(item, str) and item.strip() for item in answer), (
-            "Answer must contain only non-empty strings"
-        )
+        assert isinstance(answer, str), f"Expected string, got {type(answer).__name__}"
+        assert answer.strip(), "Answer must contain only non-empty strings"

cognee/tests/unit/modules/retrieval/graph_completion_retriever_test.py

@@ -5,6 +5,7 @@ from typing import Optional, Union
 
 import cognee
 from cognee.low_level import setup, DataPoint
+from cognee.modules.graph.utils import resolve_edges_to_text
 from cognee.tasks.storage import add_data_points
 from cognee.infrastructure.databases.exceptions import DatabaseNotCreatedError
 from cognee.modules.retrieval.graph_completion_retriever import GraphCompletionRetriever
@@ -67,7 +68,7 @@ class TestGraphCompletionRetriever:
 
         retriever = GraphCompletionRetriever()
 
-        context, _ = await retriever.get_context("Who works at Canva?")
+        context = await resolve_edges_to_text(await retriever.get_context("Who works at Canva?"))
 
         # Ensure the top-level sections are present
         assert "Nodes:" in context, "Missing 'Nodes:' section in context"
@@ -191,7 +192,7 @@ class TestGraphCompletionRetriever:
 
         retriever = GraphCompletionRetriever(top_k=20)
 
-        context, _ = await retriever.get_context("Who works at Figma?")
+        context = await resolve_edges_to_text(await retriever.get_context("Who works at Figma?"))
 
         print(context)
 
@@ -222,5 +223,5 @@ class TestGraphCompletionRetriever:
 
         await setup()
 
-        context, _ = await retriever.get_context("Who works at Figma?")
-        assert context == "", "Context should be empty on an empty graph"
+        context = await retriever.get_context("Who works at Figma?")
+        assert context == [], "Context should be empty on an empty graph"

cognee/tests/unit/modules/retrieval/insights_retriever_test.py

@@ -82,7 +82,7 @@ class TestInsightsRetriever:
 
         context = await retriever.get_context("Mike")
 
-        assert context[0][0]["name"] == "Mike Broski", "Failed to get Mike Broski"
+        assert context[0].node1.attributes["name"] == "Mike Broski", "Failed to get Mike Broski"
 
     @pytest.mark.asyncio
     async def test_insights_context_complex(self):
@@ -222,7 +222,9 @@ class TestInsightsRetriever:
 
         context = await retriever.get_context("Christina")
 
-        assert context[0][0]["name"] == "Christina Mayer", "Failed to get Christina Mayer"
+        assert context[0].node1.attributes["name"] == "Christina Mayer", (
+            "Failed to get Christina Mayer"
+        )
 
     @pytest.mark.asyncio
     async def test_insights_context_on_empty_graph(self):