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

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

@@ -31,14 +31,18 @@ async def test_get_context_success(mock_vector_engine):
         "cognee.modules.retrieval.triplet_retriever.get_vector_engine",
         return_value=mock_vector_engine,
     ):
-        context = await retriever.get_context("test query")
+        objects = await retriever.get_retrieved_objects("test query")
+        context = await retriever.get_context_from_objects("test query", objects)
+        await retriever.get_completion_from_context("test query", objects, context)
 
     assert context == "Alice knows Bob\nBob works at Tech Corp"
-    mock_vector_engine.search.assert_awaited_once_with("Triplet_text", "test query", limit=5)
+    mock_vector_engine.search.assert_awaited_once_with(
+        "Triplet_text", "test query", limit=5, include_payload=True
+    )
 
 
 @pytest.mark.asyncio
-async def test_get_context_no_collection(mock_vector_engine):
+async def test_get_objects_no_collection(mock_vector_engine):
     """Test that NoDataError is raised when Triplet_text collection doesn't exist."""
     mock_vector_engine.has_collection.return_value = False
 
@@ -49,7 +53,7 @@ async def test_get_context_no_collection(mock_vector_engine):
         return_value=mock_vector_engine,
     ):
         with pytest.raises(NoDataError, match="create_triplet_embeddings"):
-            await retriever.get_context("test query")
+            await retriever.get_retrieved_objects("test query")
 
 
 @pytest.mark.asyncio
@@ -63,13 +67,13 @@ async def test_get_context_empty_results(mock_vector_engine):
         "cognee.modules.retrieval.triplet_retriever.get_vector_engine",
         return_value=mock_vector_engine,
     ):
-        context = await retriever.get_context("test query")
+        context = await retriever.get_context_from_objects("test query", [])
 
     assert context == ""
 
 
 @pytest.mark.asyncio
-async def test_get_context_collection_not_found_error(mock_vector_engine):
+async def test_get_objects_collection_not_found_error(mock_vector_engine):
     """Test that CollectionNotFoundError is converted to NoDataError."""
     mock_vector_engine.has_collection.side_effect = CollectionNotFoundError("Collection not found")
 
@@ -80,4 +84,260 @@ async def test_get_context_collection_not_found_error(mock_vector_engine):
         return_value=mock_vector_engine,
     ):
         with pytest.raises(NoDataError, match="No data found"):
-            await retriever.get_context("test query")
+            await retriever.get_retrieved_objects("test query")
+
+
+@pytest.mark.asyncio
+async def test_get_context_empty_payload_text(mock_vector_engine):
+    """Test get_context handles missing text in payload."""
+    mock_result = MagicMock()
+    mock_result.payload = {}
+
+    mock_vector_engine.search.return_value = [mock_result]
+
+    retriever = TripletRetriever()
+
+    with patch(
+        "cognee.modules.retrieval.triplet_retriever.get_vector_engine",
+        return_value=mock_vector_engine,
+    ):
+        with pytest.raises(KeyError):
+            objects = await retriever.get_retrieved_objects("test query")
+            await retriever.get_context_from_objects("test query", retrieved_objects=objects)
+
+
+@pytest.mark.asyncio
+async def test_get_context_single_triplet(mock_vector_engine):
+    """Test get_context with single triplet result."""
+    mock_result = MagicMock()
+    mock_result.payload = {"text": "Single triplet"}
+
+    mock_vector_engine.search.return_value = [mock_result]
+
+    retriever = TripletRetriever()
+
+    with patch(
+        "cognee.modules.retrieval.triplet_retriever.get_vector_engine",
+        return_value=mock_vector_engine,
+    ):
+        objects = await retriever.get_retrieved_objects("test query")
+        context = await retriever.get_context_from_objects("test query", retrieved_objects=objects)
+
+    assert context == "Single triplet"
+
+
+@pytest.mark.asyncio
+async def test_init_defaults():
+    """Test TripletRetriever initialization with defaults."""
+    retriever = TripletRetriever()
+
+    assert retriever.user_prompt_path == "context_for_question.txt"
+    assert retriever.system_prompt_path == "answer_simple_question.txt"
+    assert retriever.top_k == 5  # Default is 5
+    assert retriever.system_prompt is None
+
+
+@pytest.mark.asyncio
+async def test_init_custom_params():
+    """Test TripletRetriever initialization with custom parameters."""
+    retriever = TripletRetriever(
+        user_prompt_path="custom_user.txt",
+        system_prompt_path="custom_system.txt",
+        system_prompt="Custom prompt",
+        top_k=10,
+    )
+
+    assert retriever.user_prompt_path == "custom_user.txt"
+    assert retriever.system_prompt_path == "custom_system.txt"
+    assert retriever.system_prompt == "Custom prompt"
+    assert retriever.top_k == 10
+
+
+@pytest.mark.asyncio
+async def test_get_completion_without_context(mock_vector_engine):
+    """Test get_completion retrieves context when not provided."""
+    mock_result = MagicMock()
+    mock_result.payload = {"text": "Test triplet"}
+    mock_vector_engine.has_collection.return_value = True
+    mock_vector_engine.search.return_value = [mock_result]
+
+    retriever = TripletRetriever()
+
+    with (
+        patch(
+            "cognee.modules.retrieval.triplet_retriever.get_vector_engine",
+            return_value=mock_vector_engine,
+        ),
+        patch(
+            "cognee.modules.retrieval.triplet_retriever.generate_completion",
+            return_value="Generated answer",
+        ),
+        patch("cognee.modules.retrieval.triplet_retriever.CacheConfig") as mock_cache_config,
+    ):
+        mock_config = MagicMock()
+        mock_config.caching = False
+        mock_cache_config.return_value = mock_config
+
+        completion = await retriever.get_completion_from_context("test query", None, None)
+
+    assert isinstance(completion, list)
+    assert len(completion) == 1
+    assert completion[0] == "Generated answer"
+
+
+@pytest.mark.asyncio
+async def test_get_completion_with_provided_context(mock_vector_engine):
+    """Test get_completion uses provided context."""
+    retriever = TripletRetriever()
+
+    with (
+        patch(
+            "cognee.modules.retrieval.triplet_retriever.generate_completion",
+            return_value="Generated answer",
+        ),
+        patch("cognee.modules.retrieval.triplet_retriever.CacheConfig") as mock_cache_config,
+    ):
+        mock_config = MagicMock()
+        mock_config.caching = False
+        mock_cache_config.return_value = mock_config
+
+        completion = await retriever.get_completion_from_context(
+            "test query", None, context="Provided context"
+        )
+
+    assert isinstance(completion, list)
+    assert len(completion) == 1
+    assert completion[0] == "Generated answer"
+
+
+@pytest.mark.asyncio
+async def test_get_completion_with_session(mock_vector_engine):
+    """Test get_completion with session caching enabled."""
+    mock_result = MagicMock()
+    mock_result.payload = {"text": "Test triplet"}
+    mock_vector_engine.has_collection.return_value = True
+    mock_vector_engine.search.return_value = [mock_result]
+
+    retriever = TripletRetriever(session_id="test_session")
+
+    mock_user = MagicMock()
+    mock_user.id = "test-user-id"
+
+    with (
+        patch(
+            "cognee.modules.retrieval.triplet_retriever.get_vector_engine",
+            return_value=mock_vector_engine,
+        ),
+        patch(
+            "cognee.modules.retrieval.triplet_retriever.get_conversation_history",
+            return_value="Previous conversation",
+        ),
+        patch(
+            "cognee.modules.retrieval.triplet_retriever.summarize_text",
+            return_value="Context summary",
+        ),
+        patch(
+            "cognee.modules.retrieval.triplet_retriever.generate_completion",
+            return_value="Generated answer",
+        ),
+        patch(
+            "cognee.modules.retrieval.triplet_retriever.save_conversation_history",
+        ) as mock_save,
+        patch("cognee.modules.retrieval.triplet_retriever.CacheConfig") as mock_cache_config,
+        patch("cognee.modules.retrieval.triplet_retriever.session_user") as mock_session_user,
+    ):
+        mock_config = MagicMock()
+        mock_config.caching = True
+        mock_cache_config.return_value = mock_config
+        mock_session_user.get.return_value = mock_user
+
+        objects = await retriever.get_retrieved_objects("test query")
+        context = await retriever.get_context_from_objects("test query", retrieved_objects=objects)
+        completion = await retriever.get_completion_from_context("test query", objects, context)
+
+    assert isinstance(completion, list)
+    assert len(completion) == 1
+    assert completion[0] == "Generated answer"
+    mock_save.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_get_completion_with_session_no_user_id(mock_vector_engine):
+    """Test get_completion with session config but no user ID."""
+    mock_result = MagicMock()
+    mock_result.payload = {"text": "Test triplet"}
+    mock_vector_engine.has_collection.return_value = True
+    mock_vector_engine.search.return_value = [mock_result]
+
+    retriever = TripletRetriever()
+
+    with (
+        patch(
+            "cognee.modules.retrieval.triplet_retriever.get_vector_engine",
+            return_value=mock_vector_engine,
+        ),
+        patch(
+            "cognee.modules.retrieval.triplet_retriever.generate_completion",
+            return_value="Generated answer",
+        ),
+        patch("cognee.modules.retrieval.triplet_retriever.CacheConfig") as mock_cache_config,
+        patch("cognee.modules.retrieval.triplet_retriever.session_user") as mock_session_user,
+    ):
+        mock_config = MagicMock()
+        mock_config.caching = True
+        mock_cache_config.return_value = mock_config
+        mock_session_user.get.return_value = None  # No user
+
+        objects = await retriever.get_retrieved_objects("test query")
+        context = await retriever.get_context_from_objects("test query", retrieved_objects=objects)
+        completion = await retriever.get_completion_from_context("test query", objects, context)
+
+    assert isinstance(completion, list)
+    assert len(completion) == 1
+
+
+@pytest.mark.asyncio
+async def test_get_completion_with_response_model(mock_vector_engine):
+    """Test get_completion with custom response model."""
+    from pydantic import BaseModel
+
+    class TestModel(BaseModel):
+        answer: str
+
+    mock_result = MagicMock()
+    mock_result.payload = {"text": "Test triplet"}
+    mock_vector_engine.has_collection.return_value = True
+    mock_vector_engine.search.return_value = [mock_result]
+
+    retriever = TripletRetriever(response_model=TestModel)
+
+    with (
+        patch(
+            "cognee.modules.retrieval.triplet_retriever.get_vector_engine",
+            return_value=mock_vector_engine,
+        ),
+        patch(
+            "cognee.modules.retrieval.triplet_retriever.generate_completion",
+            return_value=TestModel(answer="Test answer"),
+        ),
+        patch("cognee.modules.retrieval.triplet_retriever.CacheConfig") as mock_cache_config,
+    ):
+        mock_config = MagicMock()
+        mock_config.caching = False
+        mock_cache_config.return_value = mock_config
+
+        objects = await retriever.get_retrieved_objects("test query")
+        context = await retriever.get_context_from_objects("test query", retrieved_objects=objects)
+        completion = await retriever.get_completion_from_context("test query", objects, context)
+
+    assert isinstance(completion, list)
+    assert len(completion) == 1
+    assert isinstance(completion[0], TestModel)
+
+
+@pytest.mark.asyncio
+async def test_init_none_top_k():
+    """Test TripletRetriever initialization with None top_k."""
+    retriever = TripletRetriever(top_k=None)
+
+    assert retriever.top_k == 5

cognee/tests/unit/modules/search/test_get_search_type_retriever_instance.py

@@ -0,0 +1,125 @@
+import pytest
+
+from cognee.modules.search.exceptions import UnsupportedSearchTypeError
+from cognee.modules.search.types import SearchType
+
+
+class _DummyCommunityRetriever:
+    def __init__(self, *args, **kwargs):
+        self.kwargs = kwargs
+
+    def get_completion(self, *args, **kwargs):
+        return {"kind": "completion", "init": self.kwargs, "args": args, "kwargs": kwargs}
+
+    def get_context(self, *args, **kwargs):
+        return {"kind": "context", "init": self.kwargs, "args": args, "kwargs": kwargs}
+
+
+@pytest.mark.asyncio
+async def test_feeling_lucky_delegates_to_select_search_type(monkeypatch):
+    import cognee.modules.search.methods.get_search_type_retriever_instance as mod
+    from cognee.modules.retrieval.chunks_retriever import ChunksRetriever
+
+    async def _fake_select_search_type(query_text: str):
+        assert query_text == "hello"
+        return SearchType.CHUNKS
+
+    monkeypatch.setattr(mod, "select_search_type", _fake_select_search_type)
+
+    retriever_instance = await mod.get_search_type_retriever_instance(
+        SearchType.FEELING_LUCKY, query_text="hello"
+    )
+    assert isinstance(retriever_instance, ChunksRetriever)
+
+
+@pytest.mark.asyncio
+async def test_disallowed_cypher_search_types_raise(monkeypatch):
+    import cognee.modules.search.methods.get_search_type_retriever_instance as mod
+
+    monkeypatch.setenv("ALLOW_CYPHER_QUERY", "false")
+
+    with pytest.raises(UnsupportedSearchTypeError, match="disabled"):
+        await mod.get_search_type_retriever_instance(
+            SearchType.CYPHER, query_text="MATCH (n) RETURN n"
+        )
+
+    with pytest.raises(UnsupportedSearchTypeError, match="disabled"):
+        await mod.get_search_type_retriever_instance(
+            SearchType.NATURAL_LANGUAGE, query_text="Find nodes"
+        )
+
+
+@pytest.mark.asyncio
+async def test_allowed_cypher_search_types_return_tools(monkeypatch):
+    import cognee.modules.search.methods.get_search_type_retriever_instance as mod
+    from cognee.modules.retrieval.cypher_search_retriever import CypherSearchRetriever
+
+    monkeypatch.setenv("ALLOW_CYPHER_QUERY", "true")
+
+    retriever_instance = await mod.get_search_type_retriever_instance(
+        SearchType.CYPHER, query_text="q"
+    )
+    assert isinstance(retriever_instance, CypherSearchRetriever)
+
+
+@pytest.mark.asyncio
+async def test_registered_community_retriever_is_used(monkeypatch):
+    """
+    Integration point: community retrievers are loaded from the registry module and should
+    override the default mapping when present.
+    """
+    import cognee.modules.search.methods.get_search_type_retriever_instance as mod
+    from cognee.modules.retrieval import registered_community_retrievers as registry
+
+    monkeypatch.setattr(
+        registry,
+        "registered_community_retrievers",
+        {SearchType.SUMMARIES: _DummyCommunityRetriever},
+    )
+
+    retriever_instance = await mod.get_search_type_retriever_instance(
+        SearchType.SUMMARIES, query_text="q", top_k=7
+    )
+
+    assert isinstance(retriever_instance, _DummyCommunityRetriever)
+
+
+@pytest.mark.asyncio
+async def test_unknown_query_type_raises_unsupported():
+    import cognee.modules.search.methods.get_search_type_retriever_instance as mod
+
+    with pytest.raises(UnsupportedSearchTypeError, match="UNKNOWN_TYPE"):
+        await mod.get_search_type_retriever_instance("UNKNOWN_TYPE", query_text="q")
+
+
+@pytest.mark.asyncio
+async def test_default_mapping_passes_top_k_to_retrievers():
+    import cognee.modules.search.methods.get_search_type_retriever_instance as mod
+    from cognee.modules.retrieval.summaries_retriever import SummariesRetriever
+
+    retriever_instance = await mod.get_search_type_retriever_instance(
+        SearchType.SUMMARIES, query_text="q", top_k=4
+    )
+    assert isinstance(retriever_instance, SummariesRetriever)
+
+
+@pytest.mark.asyncio
+async def test_chunks_lexical_returns_jaccard_tools():
+    import cognee.modules.search.methods.get_search_type_retriever_instance as mod
+    from cognee.modules.retrieval.jaccard_retrival import JaccardChunksRetriever
+
+    retriever_instance = await mod.get_search_type_retriever_instance(
+        SearchType.CHUNKS_LEXICAL, query_text="q", top_k=3
+    )
+    assert isinstance(retriever_instance, JaccardChunksRetriever)
+
+
+@pytest.mark.asyncio
+async def test_coding_rules_uses_node_name_as_rules_nodeset_name():
+    import cognee.modules.search.methods.get_search_type_retriever_instance as mod
+    from cognee.modules.retrieval.coding_rules_retriever import CodingRulesRetriever
+
+    retriever_instance = await mod.get_search_type_retriever_instance(
+        SearchType.CODING_RULES, query_text="q", node_name=[]
+    )
+    assert isinstance(retriever_instance, CodingRulesRetriever)

cognee/tests/unit/modules/search/test_search.py

@@ -1,8 +1,9 @@
 import types
-from uuid import uuid4
+from uuid import uuid4, uuid5, NAMESPACE_OID
 
 import pytest
 
+from cognee.modules.search.models.SearchResultPayload import SearchResultPayload
 from cognee.modules.search.types import SearchType
 
 
@@ -12,9 +13,9 @@ def _make_user(user_id: str = "u1", tenant_id=None):
 
 def _make_dataset(*, name="ds", tenant_id="t1", dataset_id=None, owner_id=None):
     return types.SimpleNamespace(
-        id=dataset_id or uuid4(),
+        id=uuid5(NAMESPACE_OID, name),
         name=name,
-        tenant_id=tenant_id,
+        tenant_id=uuid5(NAMESPACE_OID, tenant_id),
         owner_id=owner_id or uuid4(),
     )
 
@@ -38,16 +39,9 @@ def _patch_side_effect_boundaries(monkeypatch, search_mod):
     async def dummy_log_result(*_args, **_kwargs):
         return None
 
-    async def dummy_prepare_search_result(search_result):
-        if isinstance(search_result, tuple) and len(search_result) == 3:
-            result, context, datasets = search_result
-            return {"result": result, "context": context, "graphs": {}, "datasets": datasets}
-        return {"result": None, "context": None, "graphs": {}, "datasets": []}
-
     monkeypatch.setattr(search_mod, "send_telemetry", lambda *a, **k: None)
     monkeypatch.setattr(search_mod, "log_query", dummy_log_query)
     monkeypatch.setattr(search_mod, "log_result", dummy_log_result)
-    monkeypatch.setattr(search_mod, "prepare_search_result", dummy_prepare_search_result)
 
     yield
 
@@ -57,9 +51,19 @@ async def test_search_access_control_returns_dataset_shaped_dicts(monkeypatch, s
     user = _make_user()
     ds = _make_dataset(name="ds1", tenant_id="t1")
 
-    async def dummy_authorized_search(**kwargs):
-        assert kwargs["dataset_ids"] == [ds.id]
-        return [("r", ["ctx"], [ds])]
+    async def dummy_authorized_search(**_kwargs):
+        assert _kwargs["dataset_ids"] == [ds.id]
+        return [
+            SearchResultPayload(
+                result_object="object",
+                context=["ctx"],
+                completion=["r"],
+                search_type=SearchType.CHUNKS,
+                dataset_name=ds.name,
+                dataset_id=ds.id,
+                dataset_tenant_id=ds.tenant_id,
+            )
+        ]
 
     monkeypatch.setattr(search_mod, "backend_access_control_enabled", lambda: True)
     monkeypatch.setattr(search_mod, "authorized_search", dummy_authorized_search)
@@ -77,24 +81,96 @@ async def test_search_access_control_returns_dataset_shaped_dicts(monkeypatch, s
             "search_result": ["r"],
             "dataset_id": ds.id,
             "dataset_name": "ds1",
-            "dataset_tenant_id": "t1",
+            "dataset_tenant_id": uuid5(NAMESPACE_OID, "t1"),
         }
     ]
 
-    out_verbose = await search_mod.search(
+
+@pytest.mark.asyncio
+async def test_search_access_control_only_context_returns_dataset_shaped_dicts(
+    monkeypatch, search_mod
+):
+    user = _make_user()
+    ds = _make_dataset(name="ds1", tenant_id="t1")
+
+    async def dummy_authorized_search(**_kwargs):
+        return [
+            SearchResultPayload(
+                result_object="object",
+                context=["ctx"],
+                completion=None,
+                search_type=SearchType.CHUNKS,
+                dataset_name=ds.name,
+                dataset_id=ds.id,
+                dataset_tenant_id=ds.tenant_id,
+            )
+        ]
+
+    monkeypatch.setattr(search_mod, "backend_access_control_enabled", lambda: True)
+    monkeypatch.setattr(search_mod, "authorized_search", dummy_authorized_search)
+
+    out = await search_mod.search(
         query_text="q",
         query_type=SearchType.CHUNKS,
         dataset_ids=[ds.id],
         user=user,
-        verbose=True,
+        only_context=True,
     )
 
-    assert out_verbose == [
+    assert out == [
         {
-            "search_result": ["r"],
+            "search_result": ["ctx"],
             "dataset_id": ds.id,
             "dataset_name": "ds1",
-            "dataset_tenant_id": "t1",
-            "graphs": {},
+            "dataset_tenant_id": uuid5(NAMESPACE_OID, "t1"),
         }
     ]
+
+
+@pytest.mark.asyncio
+async def test_authorized_search_delegates_to_search_in_datasets_context(monkeypatch, search_mod):
+    user = _make_user()
+    ds = _make_dataset(name="ds1")
+
+    async def dummy_get_authorized_existing_datasets(*_args, **_kwargs):
+        return [ds]
+
+    expected = [
+        SearchResultPayload(
+            result_object="object",
+            context="text",
+            completion="test",
+            search_type=SearchType.CHUNKS,
+            dataset_name=ds.name,
+            dataset_id=ds.id,
+            dataset_tenant_id=ds.tenant_id,
+        )
+    ]
+
+    async def dummy_search_in_datasets_context(**_kwargs):
+        return [
+            SearchResultPayload(
+                result_object="object",
+                context="text",
+                completion="test",
+                search_type=SearchType.CHUNKS,
+                dataset_name=ds.name,
+                dataset_id=ds.id,
+                dataset_tenant_id=ds.tenant_id,
+            )
+        ]
+
+    monkeypatch.setattr(
+        search_mod, "get_authorized_existing_datasets", dummy_get_authorized_existing_datasets
+    )
+    monkeypatch.setattr(search_mod, "search_in_datasets_context", dummy_search_in_datasets_context)
+
+    out = await search_mod.authorized_search(
+        query_type=SearchType.CHUNKS,
+        query_text="q",
+        user=user,
+        dataset_ids=[ds.id],
+        only_context=False,
+    )
+
+    assert out == expected