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

cognee/tests/test_search_db.py

@@ -4,6 +4,7 @@ import pathlib
 from dns.e164 import query
 
 import cognee
+from cognee.infrastructure.databases.graph import get_graph_engine
 from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge
 from cognee.modules.retrieval.graph_completion_retriever import GraphCompletionRetriever
 from cognee.modules.retrieval.graph_completion_context_extension_retriever import (
@@ -18,6 +19,7 @@ from cognee.modules.users.methods import get_default_user
 from cognee.shared.logging_utils import get_logger
 from cognee.modules.search.types import SearchType
 from cognee.modules.engine.models import NodeSet
+from collections import Counter
 
 logger = get_logger()
 
@@ -44,16 +46,16 @@ async def main():
 
     await cognee.cognify([dataset_name])
 
-    context_gk = await GraphCompletionRetriever().get_context(
+    context_gk, _ = await GraphCompletionRetriever().get_context(
         query="Next to which country is Germany located?"
     )
-    context_gk_cot = await GraphCompletionCotRetriever().get_context(
+    context_gk_cot, _ = await GraphCompletionCotRetriever().get_context(
         query="Next to which country is Germany located?"
     )
-    context_gk_ext = await GraphCompletionContextExtensionRetriever().get_context(
+    context_gk_ext, _ = await GraphCompletionContextExtensionRetriever().get_context(
         query="Next to which country is Germany located?"
     )
-    context_gk_sum = await GraphSummaryCompletionRetriever().get_context(
+    context_gk_sum, _ = await GraphSummaryCompletionRetriever().get_context(
         query="Next to which country is Germany located?"
     )
 
@@ -111,19 +113,34 @@ async def main():
 
     completion_gk = await cognee.search(
         query_type=SearchType.GRAPH_COMPLETION,
-        query_text="Next to which country is Germany located?",
+        query_text="Where is germany located, next to which country?",
+        save_interaction=True,
     )
     completion_cot = await cognee.search(
         query_type=SearchType.GRAPH_COMPLETION_COT,
-        query_text="Next to which country is Germany located?",
+        query_text="What is the country next to germany??",
+        save_interaction=True,
     )
     completion_ext = await cognee.search(
         query_type=SearchType.GRAPH_COMPLETION_CONTEXT_EXTENSION,
-        query_text="Next to which country is Germany located?",
+        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
+    )
+
     completion_sum = await cognee.search(
         query_type=SearchType.GRAPH_SUMMARY_COMPLETION,
         query_text="Next to which country is Germany located?",
+        save_interaction=True,
+    )
+
+    await cognee.search(
+        query_type=SearchType.FEEDBACK,
+        query_text="This answer was great",
+        last_k=1,
     )
 
     for name, completion in [
@@ -141,6 +158,108 @@ async def main():
             f"{name}: expected 'netherlands' in result, got: {text!r}"
         )
 
+    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])
+
+    # Assert there are exactly 4 CogneeUserInteraction nodes.
+    assert type_counts.get("CogneeUserInteraction", 0) == 4, (
+        f"Expected exactly four DCogneeUserInteraction nodes, but found {type_counts.get('CogneeUserInteraction', 0)}"
+    )
+
+    # Assert there is exactly two CogneeUserFeedback nodes.
+    assert type_counts.get("CogneeUserFeedback", 0) == 2, (
+        f"Expected exactly two CogneeUserFeedback nodes, but found {type_counts.get('CogneeUserFeedback', 0)}"
+    )
+
+    # Assert there is exactly two NodeSet.
+    assert type_counts.get("NodeSet", 0) == 2, (
+        f"Expected exactly two NodeSet nodes, but found {type_counts.get('NodeSet', 0)}"
+    )
+
+    # Assert that there are at least 10 'used_graph_element_to_answer' edges.
+    assert edge_type_counts.get("used_graph_element_to_answer", 0) >= 10, (
+        f"Expected at least ten 'used_graph_element_to_answer' edges, but found {edge_type_counts.get('used_graph_element_to_answer', 0)}"
+    )
+
+    # Assert that there are exactly 2 'gives_feedback_to' edges.
+    assert edge_type_counts.get("gives_feedback_to", 0) == 2, (
+        f"Expected exactly two 'gives_feedback_to' edges, but found {edge_type_counts.get('gives_feedback_to', 0)}"
+    )
+
+    # Assert that there are at least 6 'belongs_to_set' edges.
+    assert edge_type_counts.get("belongs_to_set", 0) == 6, (
+        f"Expected at least six 'belongs_to_set' edges, but found {edge_type_counts.get('belongs_to_set', 0)}"
+    )
+
+    nodes = graph[0]
+
+    required_fields_user_interaction = {"question", "answer", "context"}
+    required_fields_feedback = {"feedback", "sentiment"}
+
+    for node_id, data in nodes:
+        if data.get("type") == "CogneeUserInteraction":
+            assert required_fields_user_interaction.issubset(data.keys()), (
+                f"Node {node_id} is missing fields: {required_fields_user_interaction - set(data.keys())}"
+            )
+
+            for field in required_fields_user_interaction:
+                value = data[field]
+                assert isinstance(value, str) and value.strip(), (
+                    f"Node {node_id} has invalid value for '{field}': {value!r}"
+                )
+
+        if data.get("type") == "CogneeUserFeedback":
+            assert required_fields_feedback.issubset(data.keys()), (
+                f"Node {node_id} is missing fields: {required_fields_feedback - set(data.keys())}"
+            )
+
+            for field in required_fields_feedback:
+                value = data[field]
+                assert isinstance(value, str) and value.strip(), (
+                    f"Node {node_id} has invalid value for '{field}': {value!r}"
+                )
+
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+
+    await cognee.add(text_1, dataset_name)
+
+    await cognee.add([text], dataset_name)
+
+    await cognee.cognify([dataset_name])
+
+    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 = await graph_engine.get_graph_data()
+
+    edges = graph[1]
+
+    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."
+            )
+
 
 if __name__ == "__main__":
     import asyncio

cognee/tests/unit/cli/__init__.py

@@ -0,0 +1,3 @@
+"""
+CLI unit tests package.
+"""

cognee/tests/unit/cli/test_cli_commands.py

@@ -0,0 +1,483 @@
+"""
+Tests for individual CLI commands with proper mocking and coroutine handling.
+"""
+
+import pytest
+import sys
+import argparse
+import asyncio
+from unittest.mock import patch, MagicMock, AsyncMock, ANY
+from cognee.cli.commands.add_command import AddCommand
+from cognee.cli.commands.search_command import SearchCommand
+from cognee.cli.commands.cognify_command import CognifyCommand
+from cognee.cli.commands.delete_command import DeleteCommand
+from cognee.cli.commands.config_command import ConfigCommand
+from cognee.cli.exceptions import CliCommandException, CliCommandInnerException
+
+
+# Mock asyncio.run to properly handle coroutines
+def _mock_run(coro):
+    # Create an event loop and run the coroutine
+    loop = asyncio.new_event_loop()
+    try:
+        return loop.run_until_complete(coro)
+    finally:
+        loop.close()
+
+
+class TestAddCommand:
+    """Test the AddCommand class"""
+
+    def test_command_properties(self):
+        """Test basic command properties"""
+        command = AddCommand()
+        assert command.command_string == "add"
+        assert "Add data" in command.help_string
+        assert command.docs_url is not None
+
+    def test_configure_parser(self):
+        """Test parser configuration"""
+        command = AddCommand()
+        parser = argparse.ArgumentParser()
+
+        command.configure_parser(parser)
+
+        # Check that required arguments are added
+        actions = {action.dest: action for action in parser._actions}
+        assert "data" in actions
+        assert "dataset_name" in actions
+
+        # Check data argument accepts multiple values
+        assert actions["data"].nargs == "+"
+
+    @patch("cognee.cli.commands.add_command.asyncio.run", side_effect=_mock_run)
+    def test_execute_single_item(self, mock_asyncio_run):
+        """Test execute with single data item"""
+        # Mock the cognee module
+        mock_cognee = MagicMock()
+        mock_cognee.add = AsyncMock()
+
+        with patch.dict(sys.modules, {"cognee": mock_cognee}):
+            command = AddCommand()
+            args = argparse.Namespace(data=["test.txt"], dataset_name="test_dataset")
+            command.execute(args)
+
+        mock_asyncio_run.assert_called_once()
+        assert asyncio.iscoroutine(mock_asyncio_run.call_args[0][0])
+        mock_cognee.add.assert_awaited_once_with(data="test.txt", dataset_name="test_dataset")
+
+    @patch("cognee.cli.commands.add_command.asyncio.run", side_effect=_mock_run)
+    def test_execute_multiple_items(self, mock_asyncio_run):
+        """Test execute with multiple data items"""
+        # Mock the cognee module
+        mock_cognee = MagicMock()
+        mock_cognee.add = AsyncMock()
+
+        with patch.dict(sys.modules, {"cognee": mock_cognee}):
+            command = AddCommand()
+            args = argparse.Namespace(data=["test1.txt", "test2.txt"], dataset_name="test_dataset")
+            command.execute(args)
+
+        mock_asyncio_run.assert_called_once()
+        assert asyncio.iscoroutine(mock_asyncio_run.call_args[0][0])
+        mock_cognee.add.assert_awaited_once_with(
+            data=["test1.txt", "test2.txt"], dataset_name="test_dataset"
+        )
+
+    @patch("cognee.cli.commands.add_command.asyncio.run")
+    def test_execute_with_exception(self, mock_asyncio_run):
+        """Test execute handles exceptions properly"""
+        command = AddCommand()
+        args = argparse.Namespace(data=["test.txt"], dataset_name="test_dataset")
+
+        mock_asyncio_run.side_effect = Exception("Test error")
+
+        with pytest.raises(CliCommandException):
+            command.execute(args)
+
+
+class TestSearchCommand:
+    """Test the SearchCommand class"""
+
+    def test_command_properties(self):
+        """Test basic command properties"""
+        command = SearchCommand()
+        assert command.command_string == "search"
+        assert "Search and query" in command.help_string
+        assert command.docs_url is not None
+
+    def test_configure_parser(self):
+        """Test parser configuration"""
+        command = SearchCommand()
+        parser = argparse.ArgumentParser()
+
+        command.configure_parser(parser)
+
+        # Check that required arguments are added
+        actions = {action.dest: action for action in parser._actions}
+        assert "query_text" in actions
+        assert "query_type" in actions
+        assert "datasets" in actions
+        assert "top_k" in actions
+        assert "output_format" in actions
+
+        # Check default values
+        assert actions["query_type"].default == "GRAPH_COMPLETION"
+        assert actions["top_k"].default == 10
+        assert actions["output_format"].default == "pretty"
+
+    @patch("cognee.cli.commands.search_command.asyncio.run", side_effect=_mock_run)
+    def test_execute_basic_search(self, mock_asyncio_run):
+        """Test execute with basic search"""
+        # Mock the cognee module and SearchType
+        mock_cognee = MagicMock()
+        mock_cognee.search = AsyncMock(return_value=["result1", "result2"])
+        mock_search_type = MagicMock()
+        mock_search_type.__getitem__.return_value = "GRAPH_COMPLETION"
+
+        with patch.dict(sys.modules, {"cognee": mock_cognee}):
+            command = SearchCommand()
+            args = argparse.Namespace(
+                query_text="test query",
+                query_type="GRAPH_COMPLETION",
+                datasets=None,
+                top_k=10,
+                system_prompt=None,
+                output_format="pretty",
+            )
+            command.execute(args)
+
+        mock_asyncio_run.assert_called_once()
+        assert asyncio.iscoroutine(mock_asyncio_run.call_args[0][0])
+        mock_cognee.search.assert_awaited_once_with(
+            query_text="test query",
+            query_type=ANY,
+            datasets=None,
+            top_k=10,
+            system_prompt_path="answer_simple_question.txt",
+        )
+        # verify the enum’s name separately
+        called_enum = mock_cognee.search.await_args.kwargs["query_type"]
+        assert called_enum.name == "GRAPH_COMPLETION"
+
+    @patch("cognee.cli.commands.search_command.asyncio.run")
+    def test_execute_with_exception(self, mock_asyncio_run):
+        """Test execute handles exceptions properly"""
+        command = SearchCommand()
+        args = argparse.Namespace(
+            query_text="test query",
+            query_type="GRAPH_COMPLETION",
+            datasets=None,
+            top_k=10,
+            system_prompt=None,
+            output_format="pretty",
+        )
+
+        mock_asyncio_run.side_effect = Exception("Search error")
+
+        with pytest.raises(CliCommandException):
+            command.execute(args)
+
+
+class TestCognifyCommand:
+    """Test the CognifyCommand class"""
+
+    def test_command_properties(self):
+        """Test basic command properties"""
+        command = CognifyCommand()
+        assert command.command_string == "cognify"
+        assert "Transform ingested data" in command.help_string
+        assert command.docs_url is not None
+
+    def test_configure_parser(self):
+        """Test parser configuration"""
+        command = CognifyCommand()
+        parser = argparse.ArgumentParser()
+
+        command.configure_parser(parser)
+
+        # Check that arguments are added
+        actions = {action.dest: action for action in parser._actions}
+        assert "datasets" in actions
+        assert "chunk_size" in actions
+        assert "ontology_file" in actions
+        assert "chunker" in actions
+        assert "background" in actions
+        assert "verbose" in actions
+
+        # Check default values
+        assert actions["chunker"].default == "TextChunker"
+
+    @patch("cognee.cli.commands.cognify_command.asyncio.run", side_effect=_mock_run)
+    def test_execute_basic_cognify(self, mock_asyncio_run):
+        """Test execute with basic cognify"""
+        # Mock the cognee module
+        mock_cognee = MagicMock()
+        mock_cognee.cognify = AsyncMock(return_value="success")
+
+        with patch.dict(sys.modules, {"cognee": mock_cognee}):
+            command = CognifyCommand()
+            args = argparse.Namespace(
+                datasets=None,
+                chunk_size=None,
+                ontology_file=None,
+                chunker="TextChunker",
+                background=False,
+                verbose=False,
+            )
+            command.execute(args)
+
+        mock_asyncio_run.assert_called_once()
+        assert asyncio.iscoroutine(mock_asyncio_run.call_args[0][0])
+        from cognee.modules.chunking.TextChunker import TextChunker
+
+        mock_cognee.cognify.assert_awaited_once_with(
+            datasets=None,
+            chunk_size=None,
+            ontology_file_path=None,
+            chunker=TextChunker,
+            run_in_background=False,
+        )
+
+    @patch("cognee.cli.commands.cognify_command.asyncio.run")
+    def test_execute_with_exception(self, mock_asyncio_run):
+        """Test execute handles exceptions properly"""
+        command = CognifyCommand()
+        args = argparse.Namespace(
+            datasets=None,
+            chunk_size=None,
+            ontology_file=None,
+            chunker="TextChunker",
+            background=False,
+            verbose=False,
+        )
+
+        mock_asyncio_run.side_effect = Exception("Cognify error")
+
+        with pytest.raises(CliCommandException):
+            command.execute(args)
+
+
+class TestDeleteCommand:
+    """Test the DeleteCommand class"""
+
+    def test_command_properties(self):
+        """Test basic command properties"""
+        command = DeleteCommand()
+        assert command.command_string == "delete"
+        assert "Delete data" in command.help_string
+        assert command.docs_url is not None
+
+    def test_configure_parser(self):
+        """Test parser configuration"""
+        command = DeleteCommand()
+        parser = argparse.ArgumentParser()
+
+        command.configure_parser(parser)
+
+        # Check that arguments are added
+        actions = {action.dest: action for action in parser._actions}
+        assert "dataset_name" in actions
+        assert "user_id" in actions
+        assert "all" in actions
+        assert "force" in actions
+
+    @patch("cognee.cli.commands.delete_command.fmt.confirm")
+    @patch("cognee.cli.commands.delete_command.asyncio.run", side_effect=_mock_run)
+    def test_execute_delete_dataset_with_confirmation(self, mock_asyncio_run, mock_confirm):
+        """Test execute delete dataset with user confirmation"""
+        # Mock the cognee module
+        mock_cognee = MagicMock()
+        mock_cognee.delete = AsyncMock()
+
+        with patch.dict(sys.modules, {"cognee": mock_cognee}):
+            command = DeleteCommand()
+            args = argparse.Namespace(
+                dataset_name="test_dataset", user_id=None, all=False, force=False
+            )
+
+            mock_confirm.return_value = True
+
+            command.execute(args)
+
+        mock_confirm.assert_called_once_with(f"Delete dataset '{args.dataset_name}'?")
+        mock_asyncio_run.assert_called_once()
+        assert asyncio.iscoroutine(mock_asyncio_run.call_args[0][0])
+        mock_cognee.delete.assert_awaited_once_with(dataset_name="test_dataset", user_id=None)
+
+    @patch("cognee.cli.commands.delete_command.fmt.confirm")
+    def test_execute_delete_cancelled(self, mock_confirm):
+        """Test execute when user cancels deletion"""
+        command = DeleteCommand()
+        args = argparse.Namespace(dataset_name="test_dataset", user_id=None, all=False, force=False)
+
+        mock_confirm.return_value = False
+
+        # Should not raise exception, just return
+        command.execute(args)
+
+        mock_confirm.assert_called_once_with(f"Delete dataset '{args.dataset_name}'?")
+
+    @patch("cognee.cli.commands.delete_command.asyncio.run", side_effect=_mock_run)
+    def test_execute_delete_forced(self, mock_asyncio_run):
+        """Test execute delete with force flag"""
+        # Mock the cognee module
+        mock_cognee = MagicMock()
+        mock_cognee.delete = AsyncMock()
+
+        with patch.dict(sys.modules, {"cognee": mock_cognee}):
+            command = DeleteCommand()
+            args = argparse.Namespace(
+                dataset_name="test_dataset", user_id=None, all=False, force=True
+            )
+
+            command.execute(args)
+
+        mock_asyncio_run.assert_called_once()
+        assert asyncio.iscoroutine(mock_asyncio_run.call_args[0][0])
+        mock_cognee.delete.assert_awaited_once_with(dataset_name="test_dataset", user_id=None)
+
+    def test_execute_no_delete_target(self):
+        """Test execute when no delete target is specified"""
+        command = DeleteCommand()
+        args = argparse.Namespace(dataset_name=None, user_id=None, all=False, force=False)
+
+        # Should not raise exception, just return with error message
+        command.execute(args)
+
+    @patch("cognee.cli.commands.delete_command.asyncio.run")
+    def test_execute_with_exception(self, mock_asyncio_run):
+        """Test execute handles exceptions properly"""
+        command = DeleteCommand()
+        args = argparse.Namespace(dataset_name="test_dataset", user_id=None, all=False, force=True)
+
+        mock_asyncio_run.side_effect = Exception("Delete error")
+
+        with pytest.raises(CliCommandException):
+            command.execute(args)
+
+
+class TestConfigCommand:
+    """Test the ConfigCommand class"""
+
+    def test_command_properties(self):
+        """Test basic command properties"""
+        command = ConfigCommand()
+        assert command.command_string == "config"
+        assert "Manage cognee configuration" in command.help_string
+        assert command.docs_url is not None
+
+    def test_configure_parser(self):
+        """Test parser configuration"""
+        command = ConfigCommand()
+        parser = argparse.ArgumentParser()
+
+        command.configure_parser(parser)
+
+        # Check that subparsers are created
+        subparsers_actions = [
+            action for action in parser._actions if isinstance(action, argparse._SubParsersAction)
+        ]
+        assert len(subparsers_actions) == 1
+
+        subparsers = subparsers_actions[0]
+        assert "get" in subparsers.choices
+        assert "set" in subparsers.choices
+        assert "list" in subparsers.choices
+        assert "unset" in subparsers.choices
+        assert "reset" in subparsers.choices
+
+    def test_execute_no_action(self):
+        """Test execute when no config action is provided"""
+        command = ConfigCommand()
+        args = argparse.Namespace()
+
+        # Should not raise exception, just return with error message
+        command.execute(args)
+
+    @patch("builtins.__import__")
+    def test_execute_get_action(self, mock_import):
+        """Test execute get action"""
+        # Mock the cognee module
+        mock_cognee = MagicMock()
+        mock_cognee.config.get = MagicMock(return_value="openai")
+        mock_import.return_value = mock_cognee
+
+        command = ConfigCommand()
+        args = argparse.Namespace(config_action="get", key="llm_provider")
+
+        command.execute(args)
+
+    @patch("builtins.__import__")
+    def test_execute_set_action(self, mock_import):
+        """Test execute set action"""
+        # Mock the cognee module
+        mock_cognee = MagicMock()
+        mock_cognee.config.set = MagicMock()
+        mock_import.return_value = mock_cognee
+
+        command = ConfigCommand()
+        args = argparse.Namespace(config_action="set", key="llm_provider", value="anthropic")
+
+        command.execute(args)
+
+    @patch("builtins.__import__")
+    def test_execute_set_action_json_value(self, mock_import):
+        """Test execute set action with JSON value"""
+        # Mock the cognee module
+        mock_cognee = MagicMock()
+        mock_cognee.config.set = MagicMock()
+        mock_import.return_value = mock_cognee
+
+        command = ConfigCommand()
+        args = argparse.Namespace(config_action="set", key="chunk_size", value="1024")
+
+        command.execute(args)
+
+    def test_execute_list_action(self):
+        """Test execute list action"""
+        command = ConfigCommand()
+        args = argparse.Namespace(config_action="list")
+
+        # Should not raise exception
+        command.execute(args)
+
+    @patch("cognee.cli.commands.config_command.fmt.confirm")
+    def test_execute_unset_action(self, mock_confirm):
+        """Test execute unset action"""
+        # Mock the cognee module
+        mock_cognee = MagicMock()
+        mock_cognee.config.set_llm_provider = MagicMock()
+
+        with patch.dict(sys.modules, {"cognee": mock_cognee}):
+            command = ConfigCommand()
+            args = argparse.Namespace(config_action="unset", key="llm_provider", force=False)
+
+            mock_confirm.return_value = True
+
+            command.execute(args)
+
+        mock_confirm.assert_called_once()
+
+    @patch("cognee.cli.commands.config_command.fmt.confirm")
+    def test_execute_reset_action(self, mock_confirm):
+        """Test execute reset action"""
+        command = ConfigCommand()
+        args = argparse.Namespace(config_action="reset", force=False)
+
+        mock_confirm.return_value = True
+
+        # Should not raise exception
+        command.execute(args)
+
+        mock_confirm.assert_called_once()
+
+    def test_execute_with_exception(self):
+        """Test execute handles exceptions properly"""
+        # Test with an invalid action that will cause an exception in the main execute method
+        command = ConfigCommand()
+        args = argparse.Namespace(config_action="invalid_action")
+
+        # This should not raise CliCommandException, just handle it gracefully
+        # The config command handles unknown actions by showing an error message
+        command.execute(args)