From aba5f9ba55a829f54da2e806ecd7f4a1bebc8837 Mon Sep 17 00:00:00 2001 From: lxobr <122801072+lxobr@users.noreply.github.com> Date: Thu, 23 Oct 2025 17:09:15 +0200 Subject: [PATCH] test: add e2e feedback enrichment test --- .github/workflows/e2e_tests.yml | 28 ++++ cognee/tests/test_feedback_enrichment.py | 162 +++++++++++++++++++++++ 2 files changed, 190 insertions(+) create mode 100644 cognee/tests/test_feedback_enrichment.py diff --git a/.github/workflows/e2e_tests.yml b/.github/workflows/e2e_tests.yml index 775eb2912..70a4b56e6 100644 --- a/.github/workflows/e2e_tests.yml +++ b/.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 diff --git a/cognee/tests/test_feedback_enrichment.py b/cognee/tests/test_feedback_enrichment.py new file mode 100644 index 000000000..37bda95d5 --- /dev/null +++ b/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())