cognee 0.5.0.dev1__py3-none-any.whl → 0.5.1.dev0__py3-none-any.whl

cognee/tests/integration/retrieval/test_triplet_retriever.py

@@ -82,3 +82,38 @@ async def test_triplet_retriever_context_simple(setup_test_environment_with_trip
     context = await retriever.get_context("Alice")
 
     assert "Alice knows Bob" in context, "Failed to get Alice triplet"
+    assert isinstance(context, str), "Context should be a string"
+    assert len(context) > 0, "Context should not be empty"
+
+
+@pytest.mark.asyncio
+async def test_triplet_retriever_context_multiple_triplets(setup_test_environment_with_triplets):
+    """Integration test: verify TripletRetriever can retrieve multiple triplets."""
+    retriever = TripletRetriever(top_k=5)
+
+    context = await retriever.get_context("Bob")
+
+    assert "Alice knows Bob" in context or "Bob works at Tech Corp" in context, (
+        "Failed to get Bob-related triplets"
+    )
+
+
+@pytest.mark.asyncio
+async def test_triplet_retriever_top_k_limit(setup_test_environment_with_triplets):
+    """Integration test: verify TripletRetriever respects top_k parameter."""
+    retriever = TripletRetriever(top_k=1)
+
+    context = await retriever.get_context("Alice")
+
+    assert isinstance(context, str), "Context should be a string"
+
+
+@pytest.mark.asyncio
+async def test_triplet_retriever_context_empty(setup_test_environment_empty):
+    """Integration test: verify TripletRetriever handles empty graph correctly."""
+    await setup()
+
+    retriever = TripletRetriever()
+
+    with pytest.raises(NoDataError):
+        await retriever.get_context("Alice")

cognee/tests/test_custom_data_label.py

@@ -0,0 +1,68 @@
+import asyncio
+import cognee
+from cognee.shared.logging_utils import setup_logging, ERROR
+from cognee.api.v1.search import SearchType
+
+
+async def main():
+    # Create a clean slate for cognee -- reset data and system state
+    print("Resetting cognee data...")
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+    print("Data reset complete.\n")
+
+    # cognee knowledge graph will be created based on this text
+    text = """
+    Natural language processing (NLP) is an interdisciplinary
+    subfield of computer science and information retrieval.
+    """
+    from cognee.tasks.ingestion.data_item import DataItem
+
+    test_item = DataItem(text, "test_item")
+    # Add the text, and make it available for cognify
+    await cognee.add(test_item)
+
+    # Use LLMs and cognee to create knowledge graph
+    ret_val = await cognee.cognify()
+
+    query_text = "Tell me about NLP"
+    print(f"Searching cognee for insights with query: '{query_text}'")
+    # Query cognee for insights on the added text
+    search_results = await cognee.search(
+        query_type=SearchType.GRAPH_COMPLETION, query_text=query_text
+    )
+
+    print("Search results:")
+    # Display results
+    for result_text in search_results:
+        print(result_text)
+
+    from cognee.modules.data.methods.get_dataset_data import get_dataset_data
+
+    for pipeline in ret_val.values():
+        dataset_id = pipeline.dataset_id
+
+    dataset_data = await get_dataset_data(dataset_id=dataset_id)
+
+    from fastapi.encoders import jsonable_encoder
+
+    data = [
+        dict(
+            **jsonable_encoder(data),
+            dataset_id=dataset_id,
+        )
+        for data in dataset_data
+    ]
+
+    # Check if label is properly added and stored
+    assert data[0]["label"] == "test_item"
+
+
+if __name__ == "__main__":
+    logger = setup_logging(log_level=ERROR)
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+    try:
+        loop.run_until_complete(main())
+    finally:
+        loop.run_until_complete(loop.shutdown_asyncgens())