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

cognee/tests/test_search_db.py

@@ -9,7 +9,6 @@ import cognee
 from cognee.infrastructure.databases.graph import get_graph_engine
 from cognee.infrastructure.databases.vector import get_vector_engine
 from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge
-from cognee.modules.graph.utils import resolve_edges_to_text
 from cognee.modules.retrieval.graph_completion_retriever import GraphCompletionRetriever
 from cognee.modules.retrieval.graph_completion_context_extension_retriever import (
     GraphCompletionContextExtensionRetriever,
@@ -48,14 +47,14 @@ async def _reset_engines_and_prune() -> None:
         # Engine might not exist yet
         pass
 
-    from cognee.infrastructure.databases.graph.get_graph_engine import create_graph_engine
-    from cognee.infrastructure.databases.vector.create_vector_engine import create_vector_engine
+    from cognee.infrastructure.databases.graph.get_graph_engine import _create_graph_engine
+    from cognee.infrastructure.databases.vector.create_vector_engine import _create_vector_engine
     from cognee.infrastructure.databases.relational.create_relational_engine import (
         create_relational_engine,
     )
 
-    create_graph_engine.cache_clear()
-    create_vector_engine.cache_clear()
+    _create_graph_engine.cache_clear()
+    _create_vector_engine.cache_clear()
     create_relational_engine.cache_clear()
 
     await cognee.prune.prune_data()
@@ -128,11 +127,10 @@ async def setup_test_environment():
     return state
 
 
-async def setup_test_environment_for_feedback():
-    """Helper function to set up test environment for feedback weight calculation test."""
-    dataset_name = "test_dataset"
-    await _reset_engines_and_prune()
-    return await _seed_default_dataset(dataset_name=dataset_name)
+async def _get_retriever_context(retriever, query: str):
+    """Retrieve objects and resolve context via the retriever API."""
+    retrieved_objects = await retriever.get_retrieved_objects(query)
+    return await retriever.get_context_from_objects(query, retrieved_objects)
 
 
 @pytest_asyncio.fixture(scope="session")
@@ -149,26 +147,30 @@ async def e2e_state():
 
     vector_engine = get_vector_engine()
     collection = await vector_engine.search(
-        collection_name="Triplet_text", query_text="Test", limit=None
+        collection_name="Triplet_text",
+        query_text="Test",
+        limit=None,
     )
 
     # --- Retriever contexts ---
     query = "Next to which country is Germany located?"
 
     contexts = {
-        "graph_completion": await GraphCompletionRetriever().get_context(query=query),
-        "graph_completion_cot": await GraphCompletionCotRetriever().get_context(query=query),
-        "graph_completion_context_extension": await GraphCompletionContextExtensionRetriever().get_context(
-            query=query
+        "graph_completion": await _get_retriever_context(GraphCompletionRetriever(), query=query),
+        "graph_completion_cot": await _get_retriever_context(
+            GraphCompletionCotRetriever(), query=query
         ),
-        "graph_summary_completion": await GraphSummaryCompletionRetriever().get_context(
-            query=query
+        "graph_completion_context_extension": await _get_retriever_context(
+            GraphCompletionContextExtensionRetriever(), query=query
+        ),
+        "graph_summary_completion": await _get_retriever_context(
+            GraphSummaryCompletionRetriever(), query=query
         ),
-        "chunks": await ChunksRetriever(top_k=5).get_context(query=query),
-        "summaries": await SummariesRetriever(top_k=5).get_context(query=query),
-        "rag_completion": await CompletionRetriever(top_k=3).get_context(query=query),
-        "temporal": await TemporalRetriever(top_k=5).get_context(query=query),
-        "triplet": await TripletRetriever().get_context(query=query),
+        "chunks": await _get_retriever_context(ChunksRetriever(top_k=5), query=query),
+        "summaries": await _get_retriever_context(SummariesRetriever(top_k=5), query=query),
+        "rag_completion": await _get_retriever_context(CompletionRetriever(top_k=3), query=query),
+        "temporal": await _get_retriever_context(TemporalRetriever(top_k=5), query=query),
+        "triplet": await _get_retriever_context(TripletRetriever(), query=query),
     }
 
     # --- Retriever triplets + vector distance validation ---
@@ -188,57 +190,56 @@ async def e2e_state():
         query_type=SearchType.GRAPH_COMPLETION,
         query_text="Where is germany located, next to which country?",
         save_interaction=True,
+        verbose=True,
     )
     completion_cot = await cognee.search(
         query_type=SearchType.GRAPH_COMPLETION_COT,
         query_text="What is the country next to germany??",
         save_interaction=True,
+        verbose=True,
     )
     completion_ext = await cognee.search(
         query_type=SearchType.GRAPH_COMPLETION_CONTEXT_EXTENSION,
         query_text="What is the name of the country next to germany",
         save_interaction=True,
-    )
-
-    await cognee.search(
-        query_type=SearchType.FEEDBACK, query_text="This was not the best answer", last_k=1
+        verbose=True,
     )
 
     completion_sum = await cognee.search(
         query_type=SearchType.GRAPH_SUMMARY_COMPLETION,
         query_text="Next to which country is Germany located?",
         save_interaction=True,
+        verbose=True,
     )
     completion_triplet = await cognee.search(
         query_type=SearchType.TRIPLET_COMPLETION,
         query_text="Next to which country is Germany located?",
         save_interaction=True,
+        verbose=True,
     )
     completion_chunks = await cognee.search(
         query_type=SearchType.CHUNKS,
         query_text="Germany",
         save_interaction=False,
+        verbose=True,
     )
     completion_summaries = await cognee.search(
         query_type=SearchType.SUMMARIES,
         query_text="Germany",
         save_interaction=False,
+        verbose=True,
     )
     completion_rag = await cognee.search(
         query_type=SearchType.RAG_COMPLETION,
         query_text="Next to which country is Germany located?",
         save_interaction=False,
+        verbose=True,
     )
     completion_temporal = await cognee.search(
         query_type=SearchType.TEMPORAL,
         query_text="Next to which country is Germany located?",
         save_interaction=False,
-    )
-
-    await cognee.search(
-        query_type=SearchType.FEEDBACK,
-        query_text="This answer was great",
-        last_k=1,
+        verbose=True,
     )
 
     # Snapshot after all E2E operations above (used by assertion-only tests).
@@ -266,32 +267,6 @@ async def e2e_state():
     }
 
 
-@pytest_asyncio.fixture(scope="session")
-async def feedback_state():
-    """Feedback-weight scenario computed once (fresh environment)."""
-    await setup_test_environment_for_feedback()
-
-    await cognee.search(
-        query_type=SearchType.GRAPH_COMPLETION,
-        query_text="Next to which country is Germany located?",
-        save_interaction=True,
-    )
-    await cognee.search(
-        query_type=SearchType.FEEDBACK,
-        query_text="This was the best answer I've ever seen",
-        last_k=1,
-    )
-    await cognee.search(
-        query_type=SearchType.FEEDBACK,
-        query_text="Wow the correctness of this answer blows my mind",
-        last_k=1,
-    )
-
-    graph_engine = await get_graph_engine()
-    graph = await graph_engine.get_graph_data()
-    return {"graph_snapshot": graph}
-
-
 @pytest.mark.asyncio
 async def test_e2e_graph_vector_consistency(e2e_state):
     """Graph and vector stores contain the same triplet edges."""
@@ -310,28 +285,24 @@ async def test_e2e_retriever_contexts(e2e_state):
         "graph_summary_completion",
     ]:
         ctx = contexts[name]
-        assert isinstance(ctx, list), f"{name}: Context should be a list"
-        assert ctx, f"{name}: Context should not be empty"
-        ctx_text = await resolve_edges_to_text(ctx)
-        lower = ctx_text.lower()
-        assert "germany" in lower or "netherlands" in lower, (
-            f"{name}: Context did not contain 'germany' or 'netherlands'; got: {ctx!r}"
-        )
+        assert isinstance(ctx, str), f"{name}: Context should be a string"
+        assert ctx.strip(), f"{name}: Context should not be empty"
+        lower = ctx.lower()
+        assert "germany" in lower or "netherlands" in lower
 
     triplet_ctx = contexts["triplet"]
     assert isinstance(triplet_ctx, str), "triplet: Context should be a string"
     assert triplet_ctx.strip(), "triplet: Context should not be empty"
 
     chunks_ctx = contexts["chunks"]
-    assert isinstance(chunks_ctx, list), "chunks: Context should be a list"
-    assert chunks_ctx, "chunks: Context should not be empty"
-    chunks_text = "\n".join(str(item.get("text", "")) for item in chunks_ctx).lower()
+    assert isinstance(chunks_ctx, str), "chunks: Context should be a string"
+    assert chunks_ctx.strip(), "chunks: Context should not be empty"
+    chunks_text = chunks_ctx.lower()
     assert "germany" in chunks_text or "netherlands" in chunks_text
 
     summaries_ctx = contexts["summaries"]
-    assert isinstance(summaries_ctx, list), "summaries: Context should be a list"
-    assert summaries_ctx, "summaries: Context should not be empty"
-    assert any(str(item.get("text", "")).strip() for item in summaries_ctx)
+    assert isinstance(summaries_ctx, str), "summaries: Context should be a string"
+    assert summaries_ctx.strip(), "summaries: Context should not be empty"
 
     rag_ctx = contexts["rag_completion"]
     assert isinstance(rag_ctx, str), "rag_completion: Context should be a string"
@@ -350,11 +321,41 @@ async def test_e2e_retriever_triplets_have_vector_distances(e2e_state):
         assert triplets, f"{name}: Triplets list should not be empty"
         for edge in triplets:
             assert isinstance(edge, Edge), f"{name}: Elements should be Edge instances"
-            distance = edge.attributes.get("vector_distance")
-            node1_distance = edge.node1.attributes.get("vector_distance")
-            node2_distance = edge.node2.attributes.get("vector_distance")
-            assert isinstance(distance, float), f"{name}: vector_distance should be float"
+            vector_distances = edge.attributes.get("vector_distance")
+            assert vector_distances is not None, (
+                f"{name}: vector_distance should be set when retrievers return results"
+            )
+            assert isinstance(vector_distances, list) and vector_distances, (
+                f"{name}: vector_distance should be a non-empty list"
+            )
+            distance = vector_distances[0]
+            assert isinstance(distance, float), (
+                f"{name}: vector_distance[0] should be float, got {type(distance)}"
+            )
             assert 0 <= distance <= 1
+
+            node1_distances = edge.node1.attributes.get("vector_distance")
+            node2_distances = edge.node2.attributes.get("vector_distance")
+            assert node1_distances is not None, (
+                f"{name}: node1 vector_distance should be set when retrievers return results"
+            )
+            assert node2_distances is not None, (
+                f"{name}: node2 vector_distance should be set when retrievers return results"
+            )
+            assert isinstance(node1_distances, list) and node1_distances, (
+                f"{name}: node1 vector_distance should be a non-empty list"
+            )
+            assert isinstance(node2_distances, list) and node2_distances, (
+                f"{name}: node2 vector_distance should be a non-empty list"
+            )
+            node1_distance = node1_distances[0]
+            node2_distance = node2_distances[0]
+            assert isinstance(node1_distance, float), (
+                f"{name}: node1 vector_distance[0] should be float, got {type(node1_distance)}"
+            )
+            assert isinstance(node2_distance, float), (
+                f"{name}: node2 vector_distance[0] should be float, got {type(node2_distance)}"
+            )
             assert 0 <= node1_distance <= 1
             assert 0 <= node2_distance <= 1
 
@@ -387,13 +388,15 @@ async def test_e2e_search_results_and_wrappers(e2e_state):
             )
             assert wrapper.get("dataset_id"), f"{name}: missing dataset_id in wrapper"
             assert wrapper.get("dataset_name") == "test_dataset"
-            assert "graphs" in wrapper
-            text = wrapper["search_result"][0]
+            result_payload = wrapper.get("text_result")
         else:
-            text = search_results[0]
+            entry = search_results[0]
+            assert isinstance(entry, dict), f"{name}: expected dict entries"
+            result_payload = entry.get("text_result")
 
-        assert isinstance(text, str) and text.strip()
-        assert "netherlands" in text.lower()
+        text_blob = str(result_payload)
+        assert text_blob.strip()
+        assert "netherlands" in text_blob.lower()
 
     # Non-LLM search types: CHUNKS / SUMMARIES validate payload list + text
     for name in ["chunks", "summaries"]:
@@ -401,16 +404,20 @@ async def test_e2e_search_results_and_wrappers(e2e_state):
         assert isinstance(search_results, list), f"{name}: should return a list"
         assert search_results, f"{name}: should not be empty"
 
-        first = search_results[0]
-        assert isinstance(first, dict), f"{name}: expected dict entries"
+        entry = search_results[0]
+        assert isinstance(entry, dict), f"{name}: expected dict entries"
+
+        context_result = entry.get("context_result")
+        text_result = entry.get("text_result")
 
-        payloads = search_results
-        if "search_result" in first and "text" not in first:
-            payloads = (first.get("search_result") or [None])[0]
+        assert isinstance(context_result, str) and context_result.strip()
+        lower_context = context_result.lower()
+        assert "germany" in lower_context or "netherlands" in lower_context
 
-        assert isinstance(payloads, list) and payloads
-        assert isinstance(payloads[0], dict)
-        assert str(payloads[0].get("text", "")).strip()
+        assert isinstance(text_result, list) and text_result
+        first_text = text_result[0]
+        assert isinstance(first_text, dict)
+        assert str(first_text.get("text", "")).strip()
 
 
 @pytest.mark.asyncio
@@ -423,14 +430,11 @@ async def test_e2e_graph_side_effects_and_node_fields(e2e_state):
     edge_type_counts = Counter(edge_type[2] for edge_type in edges)
 
     assert type_counts.get("CogneeUserInteraction", 0) == 4
-    assert type_counts.get("CogneeUserFeedback", 0) == 2
-    assert type_counts.get("NodeSet", 0) == 2
+    assert type_counts.get("NodeSet", 0) == 1
     assert edge_type_counts.get("used_graph_element_to_answer", 0) >= 10
-    assert edge_type_counts.get("gives_feedback_to", 0) == 2
-    assert edge_type_counts.get("belongs_to_set", 0) >= 6
+    assert edge_type_counts.get("belongs_to_set", 0) >= 4
 
     required_fields_user_interaction = {"question", "answer", "context"}
-    required_fields_feedback = {"feedback", "sentiment"}
 
     for node_id, data in nodes:
         if data.get("type") == "CogneeUserInteraction":
@@ -438,20 +442,3 @@ async def test_e2e_graph_side_effects_and_node_fields(e2e_state):
             for field in required_fields_user_interaction:
                 value = data[field]
                 assert isinstance(value, str) and value.strip()
-
-        if data.get("type") == "CogneeUserFeedback":
-            assert required_fields_feedback.issubset(data.keys())
-            for field in required_fields_feedback:
-                value = data[field]
-                assert isinstance(value, str) and value.strip()
-
-
-@pytest.mark.asyncio
-async def test_e2e_feedback_weight_calculation(feedback_state):
-    """Positive feedback increases used_graph_element_to_answer feedback_weight."""
-    _nodes, edges = feedback_state["graph_snapshot"]
-    for _from_node, _to_node, relationship_name, properties in edges:
-        if relationship_name == "used_graph_element_to_answer":
-            assert properties["feedback_weight"] >= 6, (
-                "Feedback weight calculation is not correct, it should be more then 6."
-            )

cognee/tests/test_usage_logger_e2e.py

@@ -0,0 +1,268 @@
+import os
+import pytest
+import pytest_asyncio
+import asyncio
+from fastapi.testclient import TestClient
+
+import cognee
+from cognee.api.client import app
+from cognee.modules.users.methods import get_default_user, get_authenticated_user
+
+
+async def _reset_engines_and_prune():
+    """Reset db engine caches and prune data/system."""
+    try:
+        from cognee.infrastructure.databases.vector import get_vector_engine
+
+        vector_engine = get_vector_engine()
+        if hasattr(vector_engine, "engine") and hasattr(vector_engine.engine, "dispose"):
+            await vector_engine.engine.dispose(close=True)
+    except Exception:
+        pass
+
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+
+
+@pytest.fixture(scope="session")
+def event_loop():
+    """Use a single asyncio event loop for this test module."""
+    loop = asyncio.new_event_loop()
+    try:
+        yield loop
+    finally:
+        loop.close()
+
+
+@pytest.fixture(scope="session")
+def e2e_config():
+    """Configure environment for E2E tests."""
+    original_env = os.environ.copy()
+    os.environ["USAGE_LOGGING"] = "true"
+    os.environ["CACHE_BACKEND"] = "redis"
+    os.environ["CACHE_HOST"] = "localhost"
+    os.environ["CACHE_PORT"] = "6379"
+    yield
+    os.environ.clear()
+    os.environ.update(original_env)
+
+
+@pytest.fixture(scope="session")
+def authenticated_client(test_client):
+    """Override authentication to use default user."""
+
+    async def override_get_authenticated_user():
+        return await get_default_user()
+
+    app.dependency_overrides[get_authenticated_user] = override_get_authenticated_user
+    yield test_client
+    app.dependency_overrides.pop(get_authenticated_user, None)
+
+
+@pytest_asyncio.fixture(scope="session")
+async def test_data_setup():
+    """Set up test data: prune first, then add file and cognify."""
+    await _reset_engines_and_prune()
+
+    dataset_name = "test_e2e_dataset"
+    test_text = "Germany is located in Europe right next to the Netherlands."
+
+    await cognee.add(test_text, dataset_name)
+    await cognee.cognify([dataset_name])
+
+    yield dataset_name
+
+    await _reset_engines_and_prune()
+
+
+@pytest_asyncio.fixture
+async def mcp_data_setup():
+    """Set up test data for MCP tests: prune first, then add file and cognify."""
+    await _reset_engines_and_prune()
+
+    dataset_name = "test_mcp_dataset"
+    test_text = "Germany is located in Europe right next to the Netherlands."
+
+    await cognee.add(test_text, dataset_name)
+    await cognee.cognify([dataset_name])
+
+    yield dataset_name
+
+    await _reset_engines_and_prune()
+
+
+@pytest.fixture(scope="session")
+def test_client():
+    """TestClient instance for API calls."""
+    with TestClient(app) as client:
+        yield client
+
+
+@pytest_asyncio.fixture
+async def cache_engine(e2e_config):
+    """Get cache engine for log verification in test's event loop."""
+    from cognee.infrastructure.databases.cache.redis.RedisAdapter import RedisAdapter
+    from cognee.infrastructure.databases.cache.config import get_cache_config
+
+    config = get_cache_config()
+    if not config.usage_logging or config.cache_backend != "redis":
+        pytest.skip("Redis usage logging not configured")
+
+    engine = RedisAdapter(
+        host=config.cache_host,
+        port=config.cache_port,
+        username=config.cache_username,
+        password=config.cache_password,
+        log_key="usage_logs",
+    )
+    return engine
+
+
+@pytest.mark.asyncio
+async def test_api_endpoint_logging(e2e_config, authenticated_client, cache_engine):
+    """Test that API endpoints succeed and log to Redis."""
+    user = await get_default_user()
+    dataset_name = "test_e2e_api_dataset"
+
+    add_response = authenticated_client.post(
+        "/api/v1/add",
+        data={"datasetName": dataset_name},
+        files=[
+            (
+                "data",
+                (
+                    "test.txt",
+                    b"Germany is located in Europe right next to the Netherlands.",
+                    "text/plain",
+                ),
+            )
+        ],
+    )
+    assert add_response.status_code in [200, 201], f"Add endpoint failed: {add_response.text}"
+
+    cognify_response = authenticated_client.post(
+        "/api/v1/cognify",
+        json={"datasets": [dataset_name], "run_in_background": False},
+    )
+    assert cognify_response.status_code in [200, 201], (
+        f"Cognify endpoint failed: {cognify_response.text}"
+    )
+
+    search_response = authenticated_client.post(
+        "/api/v1/search",
+        json={"query": "Germany", "search_type": "GRAPH_COMPLETION", "datasets": [dataset_name]},
+    )
+    assert search_response.status_code == 200, f"Search endpoint failed: {search_response.text}"
+
+    logs = await cache_engine.get_usage_logs(str(user.id), limit=20)
+
+    add_logs = [log for log in logs if log.get("function_name") == "POST /v1/add"]
+    assert len(add_logs) > 0
+    assert add_logs[0]["type"] == "api_endpoint"
+    assert add_logs[0]["user_id"] == str(user.id)
+    assert add_logs[0]["success"] is True
+
+    cognify_logs = [log for log in logs if log.get("function_name") == "POST /v1/cognify"]
+    assert len(cognify_logs) > 0
+    assert cognify_logs[0]["type"] == "api_endpoint"
+    assert cognify_logs[0]["user_id"] == str(user.id)
+    assert cognify_logs[0]["success"] is True
+
+    search_logs = [log for log in logs if log.get("function_name") == "POST /v1/search"]
+    assert len(search_logs) > 0
+    assert search_logs[0]["type"] == "api_endpoint"
+    assert search_logs[0]["user_id"] == str(user.id)
+    assert search_logs[0]["success"] is True
+
+
+@pytest.mark.asyncio
+async def test_mcp_tool_logging(e2e_config, cache_engine):
+    """Test that MCP tools succeed and log to Redis."""
+    import sys
+    import importlib.util
+    from pathlib import Path
+
+    await _reset_engines_and_prune()
+
+    repo_root = Path(__file__).parent.parent.parent
+    mcp_src_path = repo_root / "cognee-mcp" / "src"
+    mcp_server_path = mcp_src_path / "server.py"
+
+    if not mcp_server_path.exists():
+        pytest.skip(f"MCP server not found at {mcp_server_path}")
+
+    if str(mcp_src_path) not in sys.path:
+        sys.path.insert(0, str(mcp_src_path))
+
+    spec = importlib.util.spec_from_file_location("mcp_server_module", mcp_server_path)
+    mcp_server_module = importlib.util.module_from_spec(spec)
+
+    import os
+
+    original_cwd = os.getcwd()
+    try:
+        os.chdir(str(mcp_src_path))
+        spec.loader.exec_module(mcp_server_module)
+    finally:
+        os.chdir(original_cwd)
+
+    if mcp_server_module.cognee_client is None:
+        cognee_client_path = mcp_src_path / "cognee_client.py"
+        if cognee_client_path.exists():
+            spec_client = importlib.util.spec_from_file_location(
+                "cognee_client", cognee_client_path
+            )
+            cognee_client_module = importlib.util.module_from_spec(spec_client)
+            spec_client.loader.exec_module(cognee_client_module)
+            CogneeClient = cognee_client_module.CogneeClient
+            mcp_server_module.cognee_client = CogneeClient()
+        else:
+            pytest.skip(f"CogneeClient not found at {cognee_client_path}")
+
+    test_text = "Germany is located in Europe right next to the Netherlands."
+    await mcp_server_module.cognify(data=test_text)
+    await asyncio.sleep(30.0)
+
+    list_result = await mcp_server_module.list_data()
+    assert list_result is not None, "List data should return results"
+
+    search_result = await mcp_server_module.search(
+        search_query="Germany", search_type="GRAPH_COMPLETION", top_k=5
+    )
+    assert search_result is not None, "Search should return results"
+
+    interaction_data = "User: What is Germany?\nAgent: Germany is a country in Europe."
+    await mcp_server_module.save_interaction(data=interaction_data)
+    await asyncio.sleep(30.0)
+
+    status_result = await mcp_server_module.cognify_status()
+    assert status_result is not None, "Cognify status should return results"
+
+    await mcp_server_module.prune()
+    await asyncio.sleep(0.5)
+
+    logs = await cache_engine.get_usage_logs("unknown", limit=50)
+    mcp_logs = [log for log in logs if log.get("type") == "mcp_tool"]
+    assert len(mcp_logs) > 0, (
+        f"Should have MCP tool logs with user_id='unknown'. Found logs: {[log.get('function_name') for log in logs[:5]]}"
+    )
+    assert len(mcp_logs) == 6
+    function_names = [log.get("function_name") for log in mcp_logs]
+    expected_tools = [
+        "MCP cognify",
+        "MCP list_data",
+        "MCP search",
+        "MCP save_interaction",
+        "MCP cognify_status",
+        "MCP prune",
+    ]
+
+    for expected_tool in expected_tools:
+        assert expected_tool in function_names, (
+            f"Should have {expected_tool} log. Found: {function_names}"
+        )
+
+    for log in mcp_logs:
+        assert log["type"] == "mcp_tool"
+        assert log["user_id"] == "unknown"
+        assert log["success"] is True