test: add e2e feedback enrichment test

.github/workflows/e2e_tests.yml

@@ -358,6 +358,34 @@ jobs:
           EMBEDDING_API_VERSION: ${{ secrets.EMBEDDING_API_VERSION }}
         run: uv run python ./cognee/tests/tasks/entity_extraction/entity_extraction_test.py
 
+  test-feedback-enrichment:
+    name: Test Feedback Enrichment
+    runs-on: ubuntu-22.04
+    steps:
+      - name: Check out repository
+        uses: actions/checkout@v4
+
+      - name: Cognee Setup
+        uses: ./.github/actions/cognee_setup
+        with:
+          python-version: '3.11.x'
+
+      - name: Dependencies already installed
+        run: echo "Dependencies already installed in setup"
+
+      - name: Run Feedback Enrichment Test
+        env:
+          ENV: 'dev'
+          LLM_MODEL: ${{ secrets.LLM_MODEL }}
+          LLM_ENDPOINT: ${{ secrets.LLM_ENDPOINT }}
+          LLM_API_KEY: ${{ secrets.LLM_API_KEY }}
+          LLM_API_VERSION: ${{ secrets.LLM_API_VERSION }}
+          EMBEDDING_MODEL: ${{ secrets.EMBEDDING_MODEL }}
+          EMBEDDING_ENDPOINT: ${{ secrets.EMBEDDING_ENDPOINT }}
+          EMBEDDING_API_KEY: ${{ secrets.EMBEDDING_API_KEY }}
+          EMBEDDING_API_VERSION: ${{ secrets.EMBEDDING_API_VERSION }}
+        run: uv run python ./cognee/tests/test_feedback_enrichment.py
+
   run_conversation_sessions_test:
     name: Conversation sessions test
     runs-on: ubuntu-latest

cognee/tests/test_feedback_enrichment.py

@@ -0,0 +1,162 @@
+"""
+End-to-end integration test for feedback enrichment feature.
+
+Tests the complete feedback enrichment pipeline:
+1. Add data and cognify
+2. Run search with save_interaction=True to create CogneeUserInteraction nodes
+3. Submit feedback to create CogneeUserFeedback nodes
+4. Run memify with feedback enrichment tasks to create FeedbackEnrichment nodes
+5. Verify all nodes and edges are properly created and linked in the graph
+"""
+
+import os
+import pathlib
+from collections import Counter
+
+import cognee
+from cognee.infrastructure.databases.graph import get_graph_engine
+from cognee.modules.pipelines.tasks.task import Task
+from cognee.modules.search.types import SearchType
+from cognee.shared.data_models import KnowledgeGraph
+from cognee.shared.logging_utils import get_logger
+from cognee.tasks.feedback.create_enrichments import create_enrichments
+from cognee.tasks.feedback.extract_feedback_interactions import (
+    extract_feedback_interactions,
+)
+from cognee.tasks.feedback.generate_improved_answers import generate_improved_answers
+from cognee.tasks.feedback.link_enrichments_to_feedback import (
+    link_enrichments_to_feedback,
+)
+from cognee.tasks.graph import extract_graph_from_data
+from cognee.tasks.storage import add_data_points
+
+logger = get_logger()
+
+
+async def main():
+    data_directory_path = str(
+        pathlib.Path(
+            os.path.join(
+                pathlib.Path(__file__).parent,
+                ".data_storage/test_feedback_enrichment",
+            )
+        ).resolve()
+    )
+    cognee_directory_path = str(
+        pathlib.Path(
+            os.path.join(
+                pathlib.Path(__file__).parent,
+                ".cognee_system/test_feedback_enrichment",
+            )
+        ).resolve()
+    )
+
+    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 = "feedback_enrichment_test"
+
+    await cognee.add("Cognee turns documents into AI memory.", dataset_name)
+    await cognee.cognify([dataset_name])
+
+    question_text = "Say something."
+    result = await cognee.search(
+        query_type=SearchType.GRAPH_COMPLETION,
+        query_text=question_text,
+        save_interaction=True,
+    )
+
+    assert len(result) > 0, "Search should return non-empty results"
+
+    feedback_text = "Negative feedback."
+    await cognee.search(
+        query_type=SearchType.FEEDBACK,
+        query_text=feedback_text,
+        last_k=2,
+    )
+
+    graph_engine = await get_graph_engine()
+    nodes_before, edges_before = await graph_engine.get_graph_data()
+
+    interaction_nodes_before = [
+        (node_id, props)
+        for node_id, props in nodes_before
+        if props.get("type") == "CogneeUserInteraction"
+    ]
+    feedback_nodes_before = [
+        (node_id, props)
+        for node_id, props in nodes_before
+        if props.get("type") == "CogneeUserFeedback"
+    ]
+
+    edge_types_before = Counter(edge[2] for edge in edges_before)
+
+    assert len(interaction_nodes_before) >= 1, (
+        f"Expected at least 1 CogneeUserInteraction node, found {len(interaction_nodes_before)}"
+    )
+    assert len(feedback_nodes_before) >= 1, (
+        f"Expected at least 1 CogneeUserFeedback node, found {len(feedback_nodes_before)}"
+    )
+    assert edge_types_before.get("gives_feedback_to", 0) >= 1, (
+        f"Expected at least 1 'gives_feedback_to' edge, found {edge_types_before.get('gives_feedback_to', 0)}"
+    )
+
+    extraction_tasks = [Task(extract_feedback_interactions, last_n=5)]
+    enrichment_tasks = [
+        Task(generate_improved_answers, top_k=20),
+        Task(create_enrichments),
+        Task(
+            extract_graph_from_data,
+            graph_model=KnowledgeGraph,
+            task_config={"batch_size": 10},
+        ),
+        Task(add_data_points, task_config={"batch_size": 10}),
+        Task(link_enrichments_to_feedback),
+    ]
+
+    await cognee.memify(
+        extraction_tasks=extraction_tasks,
+        enrichment_tasks=enrichment_tasks,
+        data=[{}],
+        dataset="feedback_enrichment_test_memify",
+    )
+
+    nodes_after, edges_after = await graph_engine.get_graph_data()
+
+    enrichment_nodes = [
+        (node_id, props)
+        for node_id, props in nodes_after
+        if props.get("type") == "FeedbackEnrichment"
+    ]
+
+    assert len(enrichment_nodes) >= 1, (
+        f"Expected at least 1 FeedbackEnrichment node, found {len(enrichment_nodes)}"
+    )
+
+    for node_id, props in enrichment_nodes:
+        assert "text" in props, f"FeedbackEnrichment node {node_id} missing 'text' property"
+
+    enrichment_node_ids = {node_id for node_id, _ in enrichment_nodes}
+    edges_with_enrichments = [
+        edge
+        for edge in edges_after
+        if edge[0] in enrichment_node_ids or edge[1] in enrichment_node_ids
+    ]
+
+    assert len(edges_with_enrichments) >= 1, (
+        f"Expected enrichment nodes to have at least 1 edge, found {len(edges_with_enrichments)}"
+    )
+
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+
+    logger.info("All feedback enrichment tests passed successfully")
+
+
+if __name__ == "__main__":
+    import asyncio
+
+    asyncio.run(main())